Merge "RMS: Pass in rollback_id in the calls to installd"
diff --git a/core/java/android/content/rollback/PackageRollbackInfo.java b/core/java/android/content/rollback/PackageRollbackInfo.java
index 1d0ab5a..7e5532c 100644
--- a/core/java/android/content/rollback/PackageRollbackInfo.java
+++ b/core/java/android/content/rollback/PackageRollbackInfo.java
@@ -108,6 +108,11 @@
     }
 
     /** @hide */
+    public void addPendingBackup(int userId) {
+        mPendingBackups.add(userId);
+    }
+
+    /** @hide */
     public IntArray getPendingBackups() {
         return mPendingBackups;
     }
@@ -154,6 +159,19 @@
     }
 
     /** @hide */
+    public void removePendingBackup(int userId) {
+        int idx = mPendingBackups.indexOf(userId);
+        if (idx != -1) {
+            mPendingBackups.remove(idx);
+        }
+    }
+
+    /** @hide */
+    public void removePendingRestoreInfo(int userId) {
+        removeRestoreInfo(getRestoreInfo(userId));
+    }
+
+    /** @hide */
     public PackageRollbackInfo(VersionedPackage packageRolledBackFrom,
             VersionedPackage packageRolledBackTo,
             @NonNull IntArray pendingBackups, @NonNull ArrayList<RestoreInfo> pendingRestores,
diff --git a/services/core/java/com/android/server/pm/Installer.java b/services/core/java/com/android/server/pm/Installer.java
index c2a75ab..bab612d 100644
--- a/services/core/java/com/android/server/pm/Installer.java
+++ b/services/core/java/com/android/server/pm/Installer.java
@@ -616,6 +616,7 @@
      *
      * @param pkg name of the package to snapshot user data for.
      * @param userId id of the user whose data to snapshot.
+     * @param snapshotId id of this snapshot.
      * @param storageFlags flags controlling which data (CE or DE) to snapshot.
      *
      * @return inode of the snapshot of users CE package data, or {@code 0} if a remote calls
@@ -623,12 +624,12 @@
      *
      * @throws InstallerException if failed to snapshot user data.
      */
-    public long snapshotAppData(String pkg, @UserIdInt int userId, int storageFlags)
+    public long snapshotAppData(String pkg, @UserIdInt int userId, int snapshotId, int storageFlags)
             throws InstallerException {
         if (!checkBeforeRemote()) return 0;
 
         try {
-            return mInstalld.snapshotAppData(null, pkg, userId, storageFlags);
+            return mInstalld.snapshotAppData(null, pkg, userId, snapshotId, storageFlags);
         } catch (Exception e) {
             throw InstallerException.from(e);
         }
@@ -639,8 +640,8 @@
      *
      * @param pkg name of the package to restore user data for.
      * @param appId id of the package to restore user data for.
-     * @param ceDataInode inode of CE user data folder of this app.
      * @param userId id of the user whose data to restore.
+     * @param snapshotId id of the snapshot to restore.
      * @param storageFlags flags controlling which data (CE or DE) to restore.
      *
      * @return {@code true} if user data restore was successful, or {@code false} if a remote call
@@ -648,12 +649,12 @@
      *
      * @throws InstallerException if failed to restore user data.
      */
-    public boolean restoreAppDataSnapshot(String pkg, @AppIdInt  int appId, long ceDataInode,
-            String seInfo, @UserIdInt int userId, int storageFlags) throws InstallerException {
+    public boolean restoreAppDataSnapshot(String pkg, @AppIdInt  int appId, String seInfo,
+            @UserIdInt int userId, int snapshotId, int storageFlags) throws InstallerException {
         if (!checkBeforeRemote()) return false;
 
         try {
-            mInstalld.restoreAppDataSnapshot(null, pkg, appId, ceDataInode, seInfo, userId,
+            mInstalld.restoreAppDataSnapshot(null, pkg, appId, seInfo, userId, snapshotId,
                     storageFlags);
             return true;
         } catch (Exception e) {
@@ -667,6 +668,7 @@
      * @param pkg name of the package to delete user data snapshot for.
      * @param userId id of the user whose user data snapshot to delete.
      * @param ceSnapshotInode inode of CE user data snapshot.
+     * @param snapshotId id of the snapshot to delete.
      * @param storageFlags flags controlling which user data snapshot (CE or DE) to delete.
      *
      * @return {@code true} if user data snapshot was successfully deleted, or {@code false} if a
@@ -675,11 +677,12 @@
      * @throws InstallerException if failed to delete user data snapshot.
      */
     public boolean destroyAppDataSnapshot(String pkg, @UserIdInt int userId, long ceSnapshotInode,
-            int storageFlags) throws InstallerException {
+            int snapshotId, int storageFlags) throws InstallerException {
         if (!checkBeforeRemote()) return false;
 
         try {
-            mInstalld.destroyAppDataSnapshot(null, pkg, userId, ceSnapshotInode, storageFlags);
+            mInstalld.destroyAppDataSnapshot(null, pkg, userId, ceSnapshotInode, snapshotId,
+                    storageFlags);
             return true;
         } catch (Exception e) {
             throw InstallerException.from(e);
diff --git a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
index f3b8385..e9ccea54 100644
--- a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
+++ b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
@@ -29,6 +29,8 @@
 import com.android.server.pm.Installer.InstallerException;
 
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
@@ -49,16 +51,11 @@
     }
 
     /**
-     * Creates an app data snapshot for a specified {@code packageName} for {@code installedUsers},
-     * a specified set of users for whom the package is installed.
-     *
-     * @return a {@link SnapshotAppDataResult}/
-     * @see SnapshotAppDataResult
+     * Creates an app data snapshot for a specified {@code packageRollbackInfo}. Updates said {@code
+     * packageRollbackInfo} with the inodes of the CE user data snapshot folders.
      */
-    public SnapshotAppDataResult snapshotAppData(String packageName, int[] installedUsers) {
-        final IntArray pendingBackups = new IntArray();
-        final SparseLongArray ceSnapshotInodes = new SparseLongArray();
-
+    public void snapshotAppData(int snapshotId, PackageRollbackInfo packageRollbackInfo) {
+        final int[] installedUsers = packageRollbackInfo.getInstalledUsers().toArray();
         for (int user : installedUsers) {
             final int storageFlags;
             if (isUserCredentialLocked(user)) {
@@ -66,51 +63,38 @@
                 // across app user data until the user unlocks their device.
                 Log.v(TAG, "User: " + user + " isn't unlocked, skipping CE userdata backup.");
                 storageFlags = Installer.FLAG_STORAGE_DE;
-                pendingBackups.add(user);
+                packageRollbackInfo.addPendingBackup(user);
             } else {
                 storageFlags = Installer.FLAG_STORAGE_CE | Installer.FLAG_STORAGE_DE;
             }
 
             try {
-                long ceSnapshotInode = mInstaller.snapshotAppData(packageName, user, storageFlags);
+                long ceSnapshotInode = mInstaller.snapshotAppData(
+                        packageRollbackInfo.getPackageName(), user, snapshotId, storageFlags);
                 if ((storageFlags & Installer.FLAG_STORAGE_CE) != 0) {
-                    ceSnapshotInodes.put(user, ceSnapshotInode);
+                    packageRollbackInfo.putCeSnapshotInode(user, ceSnapshotInode);
                 }
             } catch (InstallerException ie) {
-                Log.e(TAG, "Unable to create app data snapshot for: " + packageName
-                        + ", userId: " + user, ie);
+                Log.e(TAG, "Unable to create app data snapshot for: "
+                        + packageRollbackInfo.getPackageName() + ", userId: " + user, ie);
             }
         }
-
-        return new SnapshotAppDataResult(pendingBackups, ceSnapshotInodes);
     }
 
     /**
-     * Restores an app data snapshot for a specified package ({@code packageName},
-     * {@code rollbackData}) for a specified {@code userId}.
+     * Restores an app data snapshot for a specified {@code packageRollbackInfo}, for a specified
+     * {@code userId}.
      *
-     * @return {@code true} iff. a change to the {@code rollbackData} has been made. Changes to
-     *         {@code rollbackData} are restricted to the removal or addition of {@code userId} to
-     *         the list of pending backups or restores.
+     * @return {@code true} iff. a change to the {@code packageRollbackInfo} has been made. Changes
+     *         to {@code packageRollbackInfo} are restricted to the removal or addition of {@code
+     *         userId} to the list of pending backups or restores.
      */
-    public boolean restoreAppData(String packageName, RollbackData rollbackData,
-            int userId, int appId, long ceDataInode, String seInfo) {
-        if (rollbackData == null) {
-            return false;
-        }
-
-        if (!rollbackData.inProgress) {
-            Log.e(TAG, "Request to restore userData for: " + packageName
-                    + ", but no rollback in progress.");
-            return false;
-        }
-
-        PackageRollbackInfo packageInfo = RollbackManagerServiceImpl.getPackageRollbackInfo(
-                rollbackData, packageName);
+    public boolean restoreAppData(int rollbackId, PackageRollbackInfo packageRollbackInfo,
+            int userId, int appId, String seInfo) {
         int storageFlags = Installer.FLAG_STORAGE_DE;
 
-        final IntArray pendingBackups = packageInfo.getPendingBackups();
-        final List<RestoreInfo> pendingRestores = packageInfo.getPendingRestores();
+        final IntArray pendingBackups = packageRollbackInfo.getPendingBackups();
+        final List<RestoreInfo> pendingRestores = packageRollbackInfo.getPendingRestores();
         boolean changedRollbackData = false;
 
         // If we still have a userdata backup pending for this user, it implies that the user
@@ -134,53 +118,60 @@
         }
 
         try {
-            mInstaller.restoreAppDataSnapshot(packageName, appId, ceDataInode,
-                    seInfo, userId, storageFlags);
+            mInstaller.restoreAppDataSnapshot(packageRollbackInfo.getPackageName(), appId, seInfo,
+                    userId, rollbackId, storageFlags);
         } catch (InstallerException ie) {
-            Log.e(TAG, "Unable to restore app data snapshot: " + packageName, ie);
+            Log.e(TAG, "Unable to restore app data snapshot: "
+                        + packageRollbackInfo.getPackageName(), ie);
         }
 
         return changedRollbackData;
     }
 
     /**
-     * Deletes an app data data snapshot for a specified package {@code packageName} for a
-     * given {@code user}.
+     * Deletes an app data snapshot with a given {@code rollbackId} for a specified package
+     * {@code packageName} for a given {@code user}.
      */
-    public void destroyAppDataSnapshot(String packageName, int user, long ceSnapshotInode) {
+    public void destroyAppDataSnapshot(int rollbackId, PackageRollbackInfo packageRollbackInfo,
+            int user) {
         int storageFlags = Installer.FLAG_STORAGE_DE;
+        final SparseLongArray ceSnapshotInodes = packageRollbackInfo.getCeSnapshotInodes();
+        long ceSnapshotInode = ceSnapshotInodes.get(user);
         if (ceSnapshotInode > 0) {
             storageFlags |= Installer.FLAG_STORAGE_CE;
         }
         try {
-            mInstaller.destroyAppDataSnapshot(packageName, user, ceSnapshotInode, storageFlags);
+            mInstaller.destroyAppDataSnapshot(packageRollbackInfo.getPackageName(), user,
+                    ceSnapshotInode, rollbackId, storageFlags);
+            if ((storageFlags & Installer.FLAG_STORAGE_CE) != 0) {
+                ceSnapshotInodes.delete(user);
+            }
         } catch (InstallerException ie) {
-            Log.e(TAG, "Unable to delete app data snapshot for " + packageName, ie);
+            Log.e(TAG, "Unable to delete app data snapshot for "
+                        + packageRollbackInfo.getPackageName(), ie);
         }
     }
 
     /**
-     * Computes the list of pending backups and restores for {@code userId} given lists of
-     * available and recent rollbacks. Packages pending backup for the given user are added
-     * to {@code pendingBackups} and packages pending restore are added to {@code pendingRestores}
-     * along with their corresponding {@code RestoreInfo}.
+     * Computes the list of pending backups for {@code userId} given lists of available rollbacks.
+     * Packages pending backup for the given user are added to {@code pendingBackupPackages} along
+     * with their corresponding {@code PackageRollbackInfo}.
      *
-     * @return the list of {@code RollbackData} that have been modified during this computation.
+     * @return the list of {@code RollbackData} that has pending backups. Note that some of the
+     *         backups won't be performed, because they might be counteracted by pending restores.
      */
-    public List<RollbackData> computePendingBackupsAndRestores(int userId,
-            ArrayList<String> pendingBackupPackages, Map<String, RestoreInfo> pendingRestores,
-            List<RollbackData> availableRollbacks, List<RollbackInfo> recentRollbacks) {
+    private static List<RollbackData> computePendingBackups(int userId,
+            Map<String, PackageRollbackInfo> pendingBackupPackages,
+            List<RollbackData> availableRollbacks) {
         List<RollbackData> rd = new ArrayList<>();
-        // First check with the list of available rollbacks to see whether there are any
-        // pending backup operations that we've not managed to execute.
+
         for (RollbackData data : availableRollbacks) {
             for (PackageRollbackInfo info : data.packages) {
                 final IntArray pendingBackupUsers = info.getPendingBackups();
                 if (pendingBackupUsers != null) {
                     final int idx = pendingBackupUsers.indexOf(userId);
                     if (idx != -1) {
-                        pendingBackupPackages.add(info.getPackageName());
-                        pendingBackupUsers.remove(idx);
+                        pendingBackupPackages.put(info.getPackageName(), info);
                         if (rd.indexOf(data) == -1) {
                             rd.add(data);
                         }
@@ -188,23 +179,30 @@
                 }
             }
         }
+        return rd;
+    }
 
-        // Then check with the list of recently executed rollbacks to see whether there are
-        // any rollback operations
+    /**
+     * Computes the list of pending restores for {@code userId} given lists of recent rollbacks.
+     * Packages pending restore are added to {@code pendingRestores} along with their corresponding
+     * {@code PackageRollbackInfo}.
+     *
+     * @return the list of {@code RollbackInfo} that has pending restores. Note that some of the
+     *         restores won't be performed, because they might be counteracted by pending backups.
+     */
+    private static List<RollbackInfo> computePendingRestores(int userId,
+            Map<String, PackageRollbackInfo> pendingRestorePackages,
+            List<RollbackInfo> recentRollbacks) {
+        List<RollbackInfo> rd = new ArrayList<>();
+
         for (RollbackInfo data : recentRollbacks) {
             for (PackageRollbackInfo info : data.getPackages()) {
                 final RestoreInfo ri = info.getRestoreInfo(userId);
                 if (ri != null) {
-                    if (pendingBackupPackages.contains(info.getPackageName())) {
-                        // This implies that the user hasn't unlocked their device between
-                        // the request to backup data for this user and the request to restore
-                        // it, so we do nothing here.
-                        pendingBackupPackages.remove(info.getPackageName());
-                    } else {
-                        pendingRestores.put(info.getPackageName(), ri);
+                    pendingRestorePackages.put(info.getPackageName(), info);
+                    if (rd.indexOf(data) == -1) {
+                        rd.add(data);
                     }
-
-                    info.removeRestoreInfo(ri);
                 }
             }
         }
@@ -216,47 +214,77 @@
      * Commits the list of pending backups and restores for a given {@code userId}. For the pending
      * backups updates corresponding {@code changedRollbackData} with a mapping from {@code userId}
      * to a inode of theirs CE user data snapshot.
+     *
+     * @return a list {@code RollbackData} that have been changed and should be stored on disk.
      */
-    public void commitPendingBackupAndRestoreForUser(int userId,
-            ArrayList<String> pendingBackups, Map<String, RestoreInfo> pendingRestores,
-            List<RollbackData> changedRollbackData) {
-        if (!pendingBackups.isEmpty()) {
-            for (String packageName : pendingBackups) {
-                try {
-                    long ceSnapshotInode = mInstaller.snapshotAppData(packageName, userId,
-                            Installer.FLAG_STORAGE_CE);
-                    for (RollbackData data : changedRollbackData) {
-                        for (PackageRollbackInfo info : data.packages) {
-                            if (info.getPackageName().equals(packageName)) {
-                                info.putCeSnapshotInode(userId, ceSnapshotInode);
-                            }
+    public List<RollbackData> commitPendingBackupAndRestoreForUser(int userId,
+            List<RollbackData> availableRollbacks, List<RollbackInfo> recentlyExecutedRollbacks) {
+
+        final Map<String, PackageRollbackInfo> pendingBackupPackages = new HashMap<>();
+        final List<RollbackData> pendingBackups = computePendingBackups(userId,
+                pendingBackupPackages, availableRollbacks);
+
+        final Map<String, PackageRollbackInfo> pendingRestorePackages = new HashMap<>();
+        final List<RollbackInfo> pendingRestores = computePendingRestores(userId,
+                pendingRestorePackages, recentlyExecutedRollbacks);
+
+        // First remove unnecessary backups, i.e. when user did not unlock their phone between the
+        // request to backup data and the request to restore it.
+        Iterator<Map.Entry<String, PackageRollbackInfo>> iter =
+                pendingBackupPackages.entrySet().iterator();
+        while (iter.hasNext()) {
+            PackageRollbackInfo backupPackage = iter.next().getValue();
+            PackageRollbackInfo restorePackage =
+                    pendingRestorePackages.get(backupPackage.getPackageName());
+            if (restorePackage != null) {
+                backupPackage.removePendingBackup(userId);
+                backupPackage.removePendingRestoreInfo(userId);
+                iter.remove();
+                pendingRestorePackages.remove(backupPackage.getPackageName());
+            }
+        }
+
+        if (!pendingBackupPackages.isEmpty()) {
+            for (RollbackData data : pendingBackups) {
+                for (PackageRollbackInfo info : data.packages) {
+                    final IntArray pendingBackupUsers = info.getPendingBackups();
+                    final int idx = pendingBackupUsers.indexOf(userId);
+                    if (idx != -1) {
+                        try {
+                            long ceSnapshotInode = mInstaller.snapshotAppData(info.getPackageName(),
+                                    userId, data.rollbackId, Installer.FLAG_STORAGE_CE);
+                            info.putCeSnapshotInode(userId, ceSnapshotInode);
+                            pendingBackupUsers.remove(idx);
+                        } catch (InstallerException ie) {
+                            Log.e(TAG,
+                                    "Unable to create app data snapshot for: "
+                                    + info.getPackageName() + ", userId: " + userId, ie);
                         }
                     }
-                } catch (InstallerException ie) {
-                    Log.e(TAG, "Unable to create app data snapshot for: " + packageName
-                            + ", userId: " + userId, ie);
                 }
             }
         }
 
-        // TODO(narayan): Should we perform the restore before the backup for packages that have
-        // both backups and restores pending ? We could get into this case if we have a pending
-        // restore from a rollback + a snapshot request from a new restore.
-        if (!pendingRestores.isEmpty()) {
-            for (String packageName : pendingRestores.keySet()) {
-                try {
-                    final RestoreInfo ri = pendingRestores.get(packageName);
-
-                    // TODO(narayan): Verify that the user of "0" for ceDataInode is accurate
-                    // here. We know that the user has unlocked (and that their CE data is
-                    // available) so we shouldn't need to resort to the fallback path.
-                    mInstaller.restoreAppDataSnapshot(packageName, ri.appId,
-                            0 /* ceDataInode */, ri.seInfo, userId, Installer.FLAG_STORAGE_CE);
-                } catch (InstallerException ie) {
-                    Log.e(TAG, "Unable to restore app data snapshot for: " + packageName, ie);
+        if (!pendingRestorePackages.isEmpty()) {
+            for (RollbackInfo data : pendingRestores) {
+                for (PackageRollbackInfo info : data.getPackages()) {
+                    final RestoreInfo ri = info.getRestoreInfo(userId);
+                    if (ri != null) {
+                        try {
+                            mInstaller.restoreAppDataSnapshot(info.getPackageName(), ri.appId,
+                                    ri.seInfo, userId, data.getRollbackId(),
+                                    Installer.FLAG_STORAGE_CE);
+                            info.removeRestoreInfo(ri);
+                        } catch (InstallerException ie) {
+                            Log.e(TAG, "Unable to restore app data snapshot for: "
+                                    + info.getPackageName(), ie);
+                        }
+                    }
                 }
             }
         }
+
+        return pendingBackups;
     }
 
     /**
@@ -267,26 +295,4 @@
         return StorageManager.isFileEncryptedNativeOrEmulated()
                 && !StorageManager.isUserKeyUnlocked(userId);
     }
-
-    /**
-     * Encapsulates a result of {@link #snapshotAppData} method.
-     */
-    public static final class SnapshotAppDataResult {
-
-        /**
-         * A list of users for which the snapshot is pending, usually because data for one or more
-         * users is still credential locked.
-         */
-        public final IntArray pendingBackups;
-
-        /**
-         * A mapping between user and an inode of theirs CE data snapshot.
-         */
-        public final SparseLongArray ceSnapshotInodes;
-
-        public SnapshotAppDataResult(IntArray pendingBackups, SparseLongArray ceSnapshotInodes) {
-            this.pendingBackups = pendingBackups;
-            this.ceSnapshotInodes = ceSnapshotInodes;
-        }
-    }
 }
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 05d3c17..c14f126 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -32,7 +32,6 @@
 import android.content.pm.VersionedPackage;
 import android.content.rollback.IRollbackManager;
 import android.content.rollback.PackageRollbackInfo;
-import android.content.rollback.PackageRollbackInfo.RestoreInfo;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
 import android.os.Binder;
@@ -466,18 +465,17 @@
 
     void onUnlockUser(int userId) {
         getHandler().post(() -> {
-            final ArrayList<String> pendingBackupPackages = new ArrayList<>();
-            final Map<String, RestoreInfo> pendingRestorePackages = new HashMap<>();
-            final List<RollbackData> changed;
+            final List<RollbackData> availableRollbacks;
+            final List<RollbackInfo> recentlyExecutedRollbacks;
             synchronized (mLock) {
                 ensureRollbackDataLoadedLocked();
-                changed = mAppDataRollbackHelper.computePendingBackupsAndRestores(userId,
-                        pendingBackupPackages, pendingRestorePackages, mAvailableRollbacks,
-                        mRecentlyExecutedRollbacks);
+                availableRollbacks = new ArrayList<>(mAvailableRollbacks);
+                recentlyExecutedRollbacks = new ArrayList<>(mRecentlyExecutedRollbacks);
             }
 
-            mAppDataRollbackHelper.commitPendingBackupAndRestoreForUser(userId,
-                    pendingBackupPackages, pendingRestorePackages, changed);
+            final List<RollbackData> changed =
+                    mAppDataRollbackHelper.commitPendingBackupAndRestoreForUser(userId,
+                            availableRollbacks, recentlyExecutedRollbacks);
 
             for (RollbackData rd : changed) {
                 try {
@@ -844,13 +842,7 @@
             for (PackageRollbackInfo info : rd.packages) {
                 if (info.getPackageName().equals(packageName)) {
                     info.getInstalledUsers().addAll(IntArray.wrap(installedUsers));
-                    AppDataRollbackHelper.SnapshotAppDataResult rs =
-                            mAppDataRollbackHelper.snapshotAppData(packageName, installedUsers);
-                    info.getPendingBackups().addAll(rs.pendingBackups);
-                    for (int i = 0; i < rs.ceSnapshotInodes.size(); i++) {
-                        info.putCeSnapshotInode(rs.ceSnapshotInodes.keyAt(i),
-                                rs.ceSnapshotInodes.valueAt(i));
-                    }
+                    mAppDataRollbackHelper.snapshotAppData(rd.rollbackId, info);
                     try {
                         mRollbackStore.saveAvailableRollback(rd);
                     } catch (IOException ioe) {
@@ -919,18 +911,10 @@
         VersionedPackage installedVersion = new VersionedPackage(packageName,
                 pkgInfo.getLongVersionCode());
 
-        final AppDataRollbackHelper.SnapshotAppDataResult result;
-        if (snapshotUserData && !isApex) {
-            result = mAppDataRollbackHelper.snapshotAppData(packageName, installedUsers);
-        } else {
-            result = new AppDataRollbackHelper.SnapshotAppDataResult(IntArray.wrap(new int[0]),
-                new SparseLongArray());
-        }
-
         PackageRollbackInfo info = new PackageRollbackInfo(newVersion, installedVersion,
-                result.pendingBackups, new ArrayList<>(), isApex, IntArray.wrap(installedUsers),
-                result.ceSnapshotInodes);
-
+                new IntArray() /* pendingBackups */, new ArrayList<>() /* pendingRestores */,
+                isApex, IntArray.wrap(installedUsers),
+                new SparseLongArray() /* ceSnapshotInodes */);
         RollbackData data;
         try {
             int childSessionId = session.getSessionId();
@@ -948,7 +932,7 @@
                     int rollbackId = allocateRollbackIdLocked();
                     if (session.isStaged()) {
                         data = mRollbackStore.createPendingStagedRollback(rollbackId,
-                                parentSessionId);
+                            parentSessionId);
                     } else {
                         data = mRollbackStore.createAvailableRollback(rollbackId);
                     }
@@ -961,6 +945,10 @@
             return false;
         }
 
+        if (snapshotUserData && !isApex) {
+            mAppDataRollbackHelper.snapshotAppData(data.rollbackId, info);
+        }
+
         try {
             RollbackStore.backupPackageCode(data, packageName, pkgInfo.applicationInfo.sourceDir);
         } catch (IOException e) {
@@ -980,8 +968,15 @@
         getHandler().post(() -> {
             final RollbackData rollbackData = getRollbackForPackage(packageName);
             for (int userId : userIds) {
+                if (rollbackData == null || !rollbackData.inProgress) {
+                    Log.e(TAG, "Request to restore userData for: " + packageName
+                                  + ", but no rollback in progress.");
+                    continue;
+                }
+                final PackageRollbackInfo info = getPackageRollbackInfo(rollbackData, packageName);
                 final boolean changedRollbackData = mAppDataRollbackHelper.restoreAppData(
-                        packageName, rollbackData, userId, appId, ceDataInode, seInfo);
+                        rollbackData.rollbackId, info, userId, appId, seInfo);
+
                 // We've updated metadata about this rollback, so save it to flash.
                 if (changedRollbackData) {
                     try {
@@ -1252,7 +1247,7 @@
      * Returns the {@code PackageRollbackInfo} associated with {@code packageName} from
      * a specified {@code RollbackData}.
      */
-    static PackageRollbackInfo getPackageRollbackInfo(RollbackData data,
+    private static PackageRollbackInfo getPackageRollbackInfo(RollbackData data,
             String packageName) {
         for (PackageRollbackInfo info : data.packages) {
             if (info.getPackageName().equals(packageName)) {
@@ -1281,11 +1276,10 @@
     private void deleteRollback(RollbackData rollbackData) {
         for (PackageRollbackInfo info : rollbackData.packages) {
             IntArray installedUsers = info.getInstalledUsers();
-            SparseLongArray ceSnapshotInodes = info.getCeSnapshotInodes();
             for (int i = 0; i < installedUsers.size(); i++) {
                 int userId = installedUsers.get(i);
-                mAppDataRollbackHelper.destroyAppDataSnapshot(info.getPackageName(), userId,
-                        ceSnapshotInodes.get(userId, 0));
+                mAppDataRollbackHelper.destroyAppDataSnapshot(rollbackData.rollbackId, info,
+                        userId);
             }
         }
         mRollbackStore.deleteAvailableRollback(rollbackData);
diff --git a/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
index 50dbaf5..d848b2d 100644
--- a/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
+++ b/services/tests/servicestests/src/com/android/server/rollback/AppDataRollbackHelperTest.java
@@ -16,19 +16,22 @@
 
 package com.android.server.rollback;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
 import android.content.rollback.PackageRollbackInfo.RestoreInfo;
+import android.content.rollback.RollbackInfo;
 import android.util.IntArray;
 import android.util.SparseLongArray;
 
@@ -42,7 +45,9 @@
 
 import java.io.File;
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 
 @RunWith(JUnit4.class)
 public class AppDataRollbackHelperTest {
@@ -55,72 +60,54 @@
         // All users are unlocked so we should snapshot data for them.
         doReturn(true).when(helper).isUserCredentialLocked(eq(10));
         doReturn(true).when(helper).isUserCredentialLocked(eq(11));
-        AppDataRollbackHelper.SnapshotAppDataResult result = helper.snapshotAppData("com.foo.bar",
-                new int[]{10, 11});
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo.bar", new int[]{10, 11});
+        helper.snapshotAppData(5, info);
 
-        assertEquals(2, result.pendingBackups.size());
-        assertEquals(10, result.pendingBackups.get(0));
-        assertEquals(11, result.pendingBackups.get(1));
+        assertEquals(2, info.getPendingBackups().size());
+        assertEquals(10, info.getPendingBackups().get(0));
+        assertEquals(11, info.getPendingBackups().get(1));
 
-        assertEquals(0, result.ceSnapshotInodes.size());
+        assertEquals(0, info.getCeSnapshotInodes().size());
 
         InOrder inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).snapshotAppData(
-                eq("com.foo.bar"), eq(10), eq(Installer.FLAG_STORAGE_DE));
+                eq("com.foo.bar"), eq(10), eq(5), eq(Installer.FLAG_STORAGE_DE));
         inOrder.verify(installer).snapshotAppData(
-                eq("com.foo.bar"), eq(11), eq(Installer.FLAG_STORAGE_DE));
+                eq("com.foo.bar"), eq(11), eq(5), eq(Installer.FLAG_STORAGE_DE));
         inOrder.verifyNoMoreInteractions();
 
         // One of the users is unlocked but the other isn't
         doReturn(false).when(helper).isUserCredentialLocked(eq(10));
         doReturn(true).when(helper).isUserCredentialLocked(eq(11));
-        when(installer.snapshotAppData(anyString(), anyInt(), anyInt())).thenReturn(239L);
+        when(installer.snapshotAppData(anyString(), anyInt(), anyInt(), anyInt())).thenReturn(239L);
 
-        result = helper.snapshotAppData("com.foo.bar", new int[]{10, 11});
-        assertEquals(1, result.pendingBackups.size());
-        assertEquals(11, result.pendingBackups.get(0));
+        PackageRollbackInfo info2 = createPackageRollbackInfo("com.foo.bar", new int[]{10, 11});
+        helper.snapshotAppData(7, info2);
+        assertEquals(1, info2.getPendingBackups().size());
+        assertEquals(11, info2.getPendingBackups().get(0));
 
-        assertEquals(1, result.ceSnapshotInodes.size());
-        assertEquals(239L, result.ceSnapshotInodes.get(10));
+        assertEquals(1, info2.getCeSnapshotInodes().size());
+        assertEquals(239L, info2.getCeSnapshotInodes().get(10));
 
         inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).snapshotAppData(
-                eq("com.foo.bar"), eq(10),
+                eq("com.foo.bar"), eq(10), eq(7),
                 eq(Installer.FLAG_STORAGE_CE | Installer.FLAG_STORAGE_DE));
         inOrder.verify(installer).snapshotAppData(
-                eq("com.foo.bar"), eq(11), eq(Installer.FLAG_STORAGE_DE));
+                eq("com.foo.bar"), eq(11), eq(7), eq(Installer.FLAG_STORAGE_DE));
         inOrder.verifyNoMoreInteractions();
     }
 
-    private static RollbackData createInProgressRollbackData(String packageName) {
-        RollbackData data = new RollbackData(1, new File("/does/not/exist"), -1, true);
-        data.packages.add(new PackageRollbackInfo(
-                new VersionedPackage(packageName, 1), new VersionedPackage(packageName, 1),
-                new IntArray(), new ArrayList<>(), false, new IntArray(), new SparseLongArray()));
-        data.inProgress = true;
-
-        return data;
+    private static PackageRollbackInfo createPackageRollbackInfo(String packageName,
+            final int[] installedUsers) {
+        return new PackageRollbackInfo(
+                new VersionedPackage(packageName, 2), new VersionedPackage(packageName, 1),
+                new IntArray(), new ArrayList<>(), false, IntArray.wrap(installedUsers),
+                new SparseLongArray());
     }
 
-    @Test
-    public void testRestoreAppDataSnapshot_noRollbackData() throws Exception {
-        Installer installer = mock(Installer.class);
-        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
-
-        assertFalse(helper.restoreAppData("com.foo", null, 0, 0, 0, "seinfo"));
-        verifyZeroInteractions(installer);
-    }
-
-    @Test
-    public void testRestoreAppDataSnapshot_noRollbackInProgress() throws Exception {
-        Installer installer = mock(Installer.class);
-        AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
-
-        RollbackData rd = createInProgressRollbackData("com.foo");
-        // Override the in progress flag.
-        rd.inProgress = false;
-        assertFalse(helper.restoreAppData("com.foo", rd, 0, 0, 0, "seinfo"));
-        verifyZeroInteractions(installer);
+    private static PackageRollbackInfo createPackageRollbackInfo(String packageName) {
+        return createPackageRollbackInfo(packageName, new int[] {});
     }
 
     @Test
@@ -128,18 +115,20 @@
         Installer installer = mock(Installer.class);
         AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
 
-        RollbackData rd = createInProgressRollbackData("com.foo");
-        IntArray pendingBackups = rd.packages.get(0).getPendingBackups();
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo");
+        IntArray pendingBackups = info.getPendingBackups();
         pendingBackups.add(10);
         pendingBackups.add(11);
 
-        assertTrue(helper.restoreAppData("com.foo", rd, 10 /* userId */, 1, 2, "seinfo"));
+        assertTrue(helper.restoreAppData(13 /* rollbackId */, info, 10 /* userId */, 1 /* appId */,
+                      "seinfo"));
 
         // Should only require FLAG_STORAGE_DE here because we have a pending backup that we
         // didn't manage to execute.
         InOrder inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).restoreAppDataSnapshot(
-                eq("com.foo"), eq(1), eq(2L), eq("seinfo"), eq(10), eq(Installer.FLAG_STORAGE_DE));
+                eq("com.foo"), eq(1) /* appId */, eq("seinfo"), eq(10) /* userId */,
+                eq(13) /* rollbackId */, eq(Installer.FLAG_STORAGE_DE));
         inOrder.verifyNoMoreInteractions();
 
         assertEquals(1, pendingBackups.size());
@@ -152,16 +141,18 @@
         AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
         doReturn(true).when(helper).isUserCredentialLocked(eq(10));
 
-        RollbackData rd = createInProgressRollbackData("com.foo");
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo");
 
-        assertTrue(helper.restoreAppData("com.foo", rd, 10 /* userId */, 1, 2, "seinfo"));
+        assertTrue(helper.restoreAppData(73 /* rollbackId */, info, 10 /* userId */, 1 /* appId */,
+                      "seinfo"));
 
         InOrder inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).restoreAppDataSnapshot(
-                eq("com.foo"), eq(1), eq(2L), eq("seinfo"), eq(10), eq(Installer.FLAG_STORAGE_DE));
+                eq("com.foo"), eq(1) /* appId */, eq("seinfo"), eq(10) /* userId */,
+                eq(73) /* rollbackId */, eq(Installer.FLAG_STORAGE_DE));
         inOrder.verifyNoMoreInteractions();
 
-        ArrayList<RestoreInfo> pendingRestores = rd.packages.get(0).getPendingRestores();
+        ArrayList<RestoreInfo> pendingRestores = info.getPendingRestores();
         assertEquals(1, pendingRestores.size());
         assertEquals(10, pendingRestores.get(0).userId);
         assertEquals(1, pendingRestores.get(0).appId);
@@ -169,21 +160,23 @@
     }
 
     @Test
-    public void testRestoreAppDataSnapshot_availableBackupForUnockedUser() throws Exception {
+    public void testRestoreAppDataSnapshot_availableBackupForUnlockedUser() throws Exception {
         Installer installer = mock(Installer.class);
         AppDataRollbackHelper helper = spy(new AppDataRollbackHelper(installer));
         doReturn(false).when(helper).isUserCredentialLocked(eq(10));
 
-        RollbackData rd = createInProgressRollbackData("com.foo");
-        assertFalse(helper.restoreAppData("com.foo", rd, 10 /* userId */, 1, 2, "seinfo"));
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo");
+        assertFalse(helper.restoreAppData(101 /* rollbackId */, info, 10 /* userId */,
+                      1 /* appId */, "seinfo"));
 
         InOrder inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).restoreAppDataSnapshot(
-                eq("com.foo"), eq(1), eq(2L), eq("seinfo"), eq(10),
+                eq("com.foo"), eq(1) /* appId */, eq("seinfo"), eq(10) /* userId */,
+                eq(101) /* rollbackId */,
                 eq(Installer.FLAG_STORAGE_DE | Installer.FLAG_STORAGE_CE));
         inOrder.verifyNoMoreInteractions();
 
-        ArrayList<RestoreInfo> pendingRestores = rd.packages.get(0).getPendingRestores();
+        ArrayList<RestoreInfo> pendingRestores = info.getPendingRestores();
         assertEquals(0, pendingRestores.size());
     }
 
@@ -191,48 +184,99 @@
     public void destroyAppData() throws Exception {
         Installer installer = mock(Installer.class);
         AppDataRollbackHelper helper = new AppDataRollbackHelper(installer);
-        SparseLongArray ceSnapshotInodes = new SparseLongArray();
-        ceSnapshotInodes.put(11, 239L);
 
-        helper.destroyAppDataSnapshot("com.foo.bar", 10, 0L);
-        helper.destroyAppDataSnapshot("com.foo.bar", 11, 239L);
+        PackageRollbackInfo info = createPackageRollbackInfo("com.foo.bar");
+        info.putCeSnapshotInode(11, 239L);
+        helper.destroyAppDataSnapshot(5 /* rollbackId */, info, 10 /* userId */);
+        helper.destroyAppDataSnapshot(5 /* rollbackId */, info, 11 /* userId */);
 
         InOrder inOrder = Mockito.inOrder(installer);
         inOrder.verify(installer).destroyAppDataSnapshot(
-                eq("com.foo.bar"), eq(10), eq(0L),
-                eq(Installer.FLAG_STORAGE_DE));
+                eq("com.foo.bar"), eq(10) /* userId */, eq(0L) /* ceSnapshotInode */,
+                eq(5) /* rollbackId */, eq(Installer.FLAG_STORAGE_DE));
         inOrder.verify(installer).destroyAppDataSnapshot(
-                eq("com.foo.bar"), eq(11), eq(239L),
-                eq(Installer.FLAG_STORAGE_DE | Installer.FLAG_STORAGE_CE));
+                eq("com.foo.bar"), eq(11) /* userId */, eq(239L) /* ceSnapshotInode */,
+                eq(5) /* rollbackId */, eq(Installer.FLAG_STORAGE_DE | Installer.FLAG_STORAGE_CE));
         inOrder.verifyNoMoreInteractions();
+
+        assertEquals(0, info.getCeSnapshotInodes().size());
     }
 
     @Test
-    public void commitPendingBackupAndRestoreForUser_updatesRollbackData() throws Exception {
+    public void commitPendingBackupAndRestoreForUser() throws Exception {
         Installer installer = mock(Installer.class);
         AppDataRollbackHelper helper = new AppDataRollbackHelper(installer);
 
-        ArrayList<RollbackData> changedRollbackData = new ArrayList<>();
-        changedRollbackData.add(createInProgressRollbackData("com.foo.bar"));
+        when(installer.snapshotAppData(anyString(), anyInt(), anyInt(), anyInt())).thenReturn(53L);
 
-        when(installer.snapshotAppData(anyString(), anyInt(), anyInt())).thenReturn(239L);
+        // This one should be backed up.
+        PackageRollbackInfo pendingBackup = createPackageRollbackInfo("com.foo", new int[]{37, 73});
+        pendingBackup.addPendingBackup(37);
 
-        ArrayList<String> pendingBackups = new ArrayList<>();
-        pendingBackups.add("com.foo.bar");
+        // Nothing should be done for this one.
+        PackageRollbackInfo wasRecentlyRestored = createPackageRollbackInfo("com.bar",
+                new int[]{37, 73});
+        wasRecentlyRestored.addPendingBackup(37);
+        wasRecentlyRestored.getPendingRestores().add(
+                new RestoreInfo(37 /* userId */, 239 /* appId*/, "seInfo"));
 
-        helper.commitPendingBackupAndRestoreForUser(11, pendingBackups,
-                new HashMap<>() /* pendingRestores */, changedRollbackData);
+        // This one should be restored
+        PackageRollbackInfo pendingRestore = createPackageRollbackInfo("com.abc",
+                new int[]{37, 73});
+        pendingRestore.putCeSnapshotInode(37, 1543L);
+        pendingRestore.getPendingRestores().add(
+                new RestoreInfo(37 /* userId */, 57 /* appId*/, "seInfo"));
 
-        assertEquals(1, changedRollbackData.size());
-        assertEquals(1, changedRollbackData.get(0).packages.size());
-        PackageRollbackInfo info = changedRollbackData.get(0).packages.get(0);
+        // This one shouldn't be processed, because it hasn't pending backups/restores for userId
+        // 37.
+        PackageRollbackInfo ignoredInfo = createPackageRollbackInfo("com.bar",
+                new int[]{3, 73});
+        wasRecentlyRestored.addPendingBackup(3);
+        wasRecentlyRestored.addPendingBackup(73);
+        wasRecentlyRestored.getPendingRestores().add(
+                new RestoreInfo(73 /* userId */, 239 /* appId*/, "seInfo"));
 
-        assertEquals(1, info.getCeSnapshotInodes().size());
-        assertEquals(239L, info.getCeSnapshotInodes().get(11));
+        RollbackData dataWithPendingBackup = new RollbackData(101, new File("/does/not/exist"), -1,
+                true);
+        dataWithPendingBackup.packages.add(pendingBackup);
 
+        RollbackData dataWithRecentRestore = new RollbackData(17239, new File("/does/not/exist"),
+                -1, true);
+        dataWithRecentRestore.packages.add(wasRecentlyRestored);
+
+        RollbackData dataForDifferentUser = new RollbackData(17239, new File("/does/not/exist"),
+                -1, true);
+        dataForDifferentUser.packages.add(ignoredInfo);
+
+        RollbackInfo rollbackInfo = new RollbackInfo(17239,
+                Arrays.asList(pendingRestore, wasRecentlyRestored), false);
+
+        List<RollbackData> changed = helper.commitPendingBackupAndRestoreForUser(37,
+                Arrays.asList(dataWithPendingBackup, dataWithRecentRestore, dataForDifferentUser),
+                Collections.singletonList(rollbackInfo));
         InOrder inOrder = Mockito.inOrder(installer);
-        inOrder.verify(installer).snapshotAppData("com.foo.bar", 11 /* userId */,
-                Installer.FLAG_STORAGE_CE);
+
+        // Check that pending backup and restore for the same package mutually destroyed each other.
+        assertEquals(-1, wasRecentlyRestored.getPendingBackups().indexOf(37));
+        assertNull(wasRecentlyRestored.getRestoreInfo(37));
+
+        // Check that backup was performed.
+        inOrder.verify(installer).snapshotAppData(eq("com.foo"), eq(37), eq(101),
+                eq(Installer.FLAG_STORAGE_CE));
+        assertEquals(-1, pendingBackup.getPendingBackups().indexOf(37));
+        assertEquals(53, pendingBackup.getCeSnapshotInodes().get(37));
+
+        // Check that changed returns correct RollbackData.
+        assertEquals(2, changed.size());
+        assertEquals(dataWithPendingBackup, changed.get(0));
+        assertEquals(dataWithRecentRestore, changed.get(1));
+
+        // Check that restore was performed.
+        inOrder.verify(installer).restoreAppDataSnapshot(
+                eq("com.abc"), eq(57) /* appId */, eq("seInfo"), eq(37) /* userId */,
+                eq(17239) /* rollbackId */, eq(Installer.FLAG_STORAGE_CE));
+        assertNull(pendingRestore.getRestoreInfo(37));
+
         inOrder.verifyNoMoreInteractions();
     }
 }
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
index f3edf09..bd0881f 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
@@ -587,9 +587,13 @@
             RollbackTestUtils.installMultiPackage(false,
                     "RollbackTestAppAv1.apk",
                     "RollbackTestAppBv1.apk");
+            processUserData(TEST_APP_A);
+            processUserData(TEST_APP_B);
             RollbackTestUtils.installMultiPackage(true,
                     "RollbackTestAppAv2.apk",
                     "RollbackTestAppBv2.apk");
+            processUserData(TEST_APP_A);
+            processUserData(TEST_APP_B);
             assertEquals(2, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
             assertEquals(2, RollbackTestUtils.getInstalledVersion(TEST_APP_B));
 
@@ -614,6 +618,9 @@
             assertRollbackInfoForAandB(rollbackB);
 
             assertEquals(rollbackA.getRollbackId(), rollbackB.getRollbackId());
+
+            processUserData(TEST_APP_A);
+            processUserData(TEST_APP_B);
         } finally {
             RollbackTestUtils.dropShellPermissionIdentity();
         }