Add per-Rollback synchronisation.

This adds a lock object to each Rollback that should be acquired before
accessing any rollback state.

For now, this synchronisation has been added 'in-place' in
RollbackManagerServiceImpl. Future changes will move rollback functionality
from that class into the Rollback class, which will allow the synchronisation
to be handled internally within the Rollback.

Bug: 136241838
Test: atest RollbackTest
Test: atest CtsStagedInstallHostTestCases
Test: atest CtsRollbackManagerHostTestCases
Change-Id: I995f1647ccf0410a8832f097236482552729f4c8

4 files changed, 346 insertions, 234 deletions

diff --git a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
index 8c0dda298246..3f9cc83b53d8 100644
--- a/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
+++ b/services/core/java/com/android/server/rollback/AppDataRollbackHelper.java
@@ -23,6 +23,7 @@ import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseLongArray;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.pm.Installer;
 import com.android.server.pm.Installer.InstallerException;
@@ -50,6 +51,8 @@ public class AppDataRollbackHelper {
      * {@code userIds}. Updates said {@code packageRollbackInfo} with the inodes of the CE user data
      * snapshot folders.
      */
+    @GuardedBy("rollback.getLock")
+    // TODO(b/136241838): Move into Rollback and synchronize there.
     public void snapshotAppData(
             int snapshotId, PackageRollbackInfo packageRollbackInfo, int[] userIds) {
         for (int user : userIds) {
@@ -86,6 +89,8 @@ public class AppDataRollbackHelper {
      *         to {@code packageRollbackInfo} are restricted to the removal or addition of {@code
      *         userId} to the list of pending backups or restores.
      */
+    @GuardedBy("rollback.getLock")
+    // TODO(b/136241838): Move into Rollback and synchronize there.
     public boolean restoreAppData(int rollbackId, PackageRollbackInfo packageRollbackInfo,
             int userId, int appId, String seInfo) {
         int storageFlags = Installer.FLAG_STORAGE_DE;
@@ -129,6 +134,8 @@ public class AppDataRollbackHelper {
      * Deletes an app data snapshot with a given {@code rollbackId} for a specified package
      * {@code packageName} for a given {@code user}.
      */
+    @GuardedBy("rollback.getLock")
+    // TODO(b/136241838): Move into Rollback and synchronize there.
     public void destroyAppDataSnapshot(int rollbackId, PackageRollbackInfo packageRollbackInfo,
             int user) {
         int storageFlags = Installer.FLAG_STORAGE_DE;
@@ -156,6 +163,7 @@ public class AppDataRollbackHelper {
      *
      * @return true if any backups or restores were found for the userId
      */
+    @GuardedBy("rollback.getLock")
     boolean commitPendingBackupAndRestoreForUser(int userId, Rollback rollback) {
         boolean foundBackupOrRestore = false;
         for (PackageRollbackInfo info : rollback.info.getPackages()) {
diff --git a/services/core/java/com/android/server/rollback/Rollback.java b/services/core/java/com/android/server/rollback/Rollback.java
index 6769fe07bbf8..80b18ea1b92e 100644
--- a/services/core/java/com/android/server/rollback/Rollback.java
+++ b/services/core/java/com/android/server/rollback/Rollback.java
@@ -18,8 +18,11 @@ package com.android.server.rollback;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
+import android.content.rollback.PackageRollbackInfo;
 import android.content.rollback.RollbackInfo;
 
+import com.android.internal.annotations.GuardedBy;
+
 import java.io.File;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -29,8 +32,12 @@ import java.util.ArrayList;
 
 
 /**
- * Information about a rollback available for a set of atomically installed
- * packages.
+ * Information about a rollback available for a set of atomically installed packages.
+ *
+ * <p>When accessing the state of a Rollback object, the caller is responsible for synchronization.
+ * The lock object provided by {@link #getLock} should be acquired when accessing any of the mutable
+ * state of a Rollback, including from the {@link RollbackInfo} and any of the
+ * {@link PackageRollbackInfo} objects held within.
  */
 class Rollback {
     @IntDef(flag = true, prefix = { "ROLLBACK_STATE_" }, value = {
@@ -58,8 +65,18 @@ class Rollback {
     static final int ROLLBACK_STATE_COMMITTED = 3;
 
     /**
+     * The session ID for the staged session if this rollback data represents a staged session,
+     * {@code -1} otherwise.
+     */
+    private final int mStagedSessionId;
+
+    /**
      * The rollback info for this rollback.
+     *
+     * <p>Any access to this field that touches any mutable state should be synchronized on
+     * {@link #getLock}.
      */
+    @GuardedBy("getLock")
     public final RollbackInfo info;
 
     /**
@@ -74,23 +91,20 @@ class Rollback {
      * The timestamp is not applicable for all rollback states, but we make
      * sure to keep it non-null to avoid potential errors there.
      */
+    @GuardedBy("mLock")
     private @NonNull Instant mTimestamp;
 
     /**
-     * The session ID for the staged session if this rollback data represents a staged session,
-     * {@code -1} otherwise.
-     */
-    private final int mStagedSessionId;
-
-    /**
      * The current state of the rollback.
      * ENABLING, AVAILABLE, or COMMITTED.
      */
+    @GuardedBy("mLock")
     private @RollbackState int mState;
 
     /**
      * The id of the post-reboot apk session for a staged install, if any.
      */
+    @GuardedBy("mLock")
     private int mApkSessionId = -1;
 
     /**
@@ -98,10 +112,17 @@ class Rollback {
      * for this rollback because it has just been committed but the rollback
      * has not yet been fully applied.
      */
-    // NOTE: All accesses to this field are from the RollbackManager handler thread.
+    @GuardedBy("mLock")
     private boolean mRestoreUserDataInProgress = false;
 
     /**
+     * Lock object to guard all access to Rollback state.
+     *
+     * @see #getLock
+     */
+    private final Object mLock = new Object();
+
+    /**
      * Constructs a new, empty Rollback instance.
      *
      * @param rollbackId the id of the rollback.
@@ -135,8 +156,23 @@ class Rollback {
     }
 
     /**
+     * Returns a lock object that should be acquired before accessing any Rollback state from
+     * {@link RollbackManagerServiceImpl}.
+     *
+     * <p>Note that while holding this lock, the lock for {@link RollbackManagerServiceImpl} should
+     * not be acquired (but it is ok to acquire this lock while already holding the lock for that
+     * class).
+     */
+    // TODO(b/136241838): Move rollback functionality into this class and synchronize on the lock
+    // internally. Remove this method once this has been done for all cases.
+    Object getLock() {
+        return mLock;
+    }
+
+    /**
      * Whether the rollback is for rollback of a staged install.
      */
+    @GuardedBy("getLock")
     boolean isStaged() {
         return info.isStaged();
     }
@@ -151,6 +187,7 @@ class Rollback {
     /**
      * Returns the time when the upgrade occurred, for purposes of expiring rollback data.
      */
+    @GuardedBy("getLock")
     Instant getTimestamp() {
         return mTimestamp;
     }
@@ -158,6 +195,7 @@ class Rollback {
     /**
      * Sets the time at which upgrade occurred.
      */
+    @GuardedBy("getLock")
     void setTimestamp(Instant timestamp) {
         mTimestamp = timestamp;
     }
@@ -173,6 +211,7 @@ class Rollback {
     /**
      * Returns true if the rollback is in the ENABLING state.
      */
+    @GuardedBy("getLock")
     boolean isEnabling() {
         return mState == ROLLBACK_STATE_ENABLING;
     }
@@ -180,6 +219,7 @@ class Rollback {
     /**
      * Returns true if the rollback is in the AVAILABLE state.
      */
+    @GuardedBy("getLock")
     boolean isAvailable() {
         return mState == ROLLBACK_STATE_AVAILABLE;
     }
@@ -187,6 +227,7 @@ class Rollback {
     /**
      * Returns true if the rollback is in the COMMITTED state.
      */
+    @GuardedBy("getLock")
     boolean isCommitted() {
         return mState == ROLLBACK_STATE_COMMITTED;
     }
@@ -194,6 +235,7 @@ class Rollback {
     /**
      * Sets the state of the rollback to AVAILABLE.
      */
+    @GuardedBy("getLock")
     void setAvailable() {
         mState = ROLLBACK_STATE_AVAILABLE;
     }
@@ -201,6 +243,7 @@ class Rollback {
     /**
      * Sets the state of the rollback to COMMITTED.
      */
+    @GuardedBy("getLock")
     void setCommitted() {
         mState = ROLLBACK_STATE_COMMITTED;
     }
@@ -208,6 +251,7 @@ class Rollback {
     /**
      * Returns the id of the post-reboot apk session for a staged install, if any.
      */
+    @GuardedBy("getLock")
     int getApkSessionId() {
         return mApkSessionId;
     }
@@ -215,6 +259,7 @@ class Rollback {
     /**
      * Sets the id of the post-reboot apk session for a staged install.
      */
+    @GuardedBy("getLock")
     void setApkSessionId(int apkSessionId) {
         mApkSessionId = apkSessionId;
     }
@@ -223,6 +268,7 @@ class Rollback {
      * Returns true if we are expecting the package manager to call restoreUserData for this
      * rollback because it has just been committed but the rollback has not yet been fully applied.
      */
+    @GuardedBy("getLock")
     boolean isRestoreUserDataInProgress() {
         return mRestoreUserDataInProgress;
     }
@@ -231,6 +277,7 @@ class Rollback {
      * Sets whether we are expecting the package manager to call restoreUserData for this
      * rollback because it has just been committed but the rollback has not yet been fully applied.
      */
+    @GuardedBy("getLock")
     void setRestoreUserDataInProgress(boolean restoreUserDataInProgress) {
         mRestoreUserDataInProgress = restoreUserDataInProgress;
     }
@@ -254,6 +301,7 @@ class Rollback {
         throw new ParseException("Invalid rollback state: " + state, 0);
     }
 
+    @GuardedBy("getLock")
     String getStateAsString() {
         return rollbackStateToString(mState);
     }
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 8b0f79a84429..12a5c8bff03e 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -282,8 +282,10 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             List<RollbackInfo> rollbacks = new ArrayList<>();
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.isAvailable()) {
-                    rollbacks.add(rollback.info);
+                synchronized (rollback.getLock()) {
+                    if (rollback.isAvailable()) {
+                        rollbacks.add(rollback.info);
+                    }
                 }
             }
             return new ParceledListSlice<>(rollbacks);
@@ -298,8 +300,10 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             List<RollbackInfo> rollbacks = new ArrayList<>();
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.isCommitted()) {
-                    rollbacks.add(rollback.info);
+                synchronized (rollback.getLock()) {
+                    if (rollback.isCommitted()) {
+                        rollbacks.add(rollback.info);
+                    }
                 }
             }
             return new ParceledListSlice<>(rollbacks);
@@ -332,8 +336,11 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                     Iterator<Rollback> iter = mRollbacks.iterator();
                     while (iter.hasNext()) {
                         Rollback rollback = iter.next();
-                        rollback.setTimestamp(rollback.getTimestamp().plusMillis(timeDifference));
-                        saveRollback(rollback);
+                        synchronized (rollback.getLock()) {
+                            rollback.setTimestamp(
+                                    rollback.getTimestamp().plusMillis(timeDifference));
+                            saveRollback(rollback);
+                        }
                     }
                 }
             }
@@ -358,86 +365,94 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         Slog.i(TAG, "Initiating rollback");
 
         Rollback rollback = getRollbackForId(rollbackId);
-        if (rollback == null || !rollback.isAvailable()) {
+        if (rollback == null) {
             sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE_ROLLBACK_UNAVAILABLE,
                     "Rollback unavailable");
             return;
         }
-
-        // Get a context for the caller to use to install the downgraded
-        // version of the package.
-        final Context context;
-        try {
-            context = mContext.createPackageContext(callerPackageName, 0);
-        } catch (PackageManager.NameNotFoundException e) {
-            sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE,
-                    "Invalid callerPackageName");
-            return;
-        }
-
-        PackageManager pm = context.getPackageManager();
-        try {
-            PackageInstaller packageInstaller = pm.getPackageInstaller();
-            PackageInstaller.SessionParams parentParams = new PackageInstaller.SessionParams(
-                    PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-            parentParams.setRequestDowngrade(true);
-            parentParams.setMultiPackage();
-            if (rollback.isStaged()) {
-                parentParams.setStaged();
+        synchronized (rollback.getLock()) {
+            if (!rollback.isAvailable()) {
+                sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE_ROLLBACK_UNAVAILABLE,
+                        "Rollback unavailable");
+                return;
             }
 
-            int parentSessionId = packageInstaller.createSession(parentParams);
-            PackageInstaller.Session parentSession = packageInstaller.openSession(parentSessionId);
+            // Get a context for the caller to use to install the downgraded
+            // version of the package.
+            final Context context;
+            try {
+                context = mContext.createPackageContext(callerPackageName, 0);
+            } catch (PackageManager.NameNotFoundException e) {
+                sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE,
+                        "Invalid callerPackageName");
+                return;
+            }
 
-            for (PackageRollbackInfo info : rollback.info.getPackages()) {
-                PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
+            PackageManager pm = context.getPackageManager();
+            try {
+                PackageInstaller packageInstaller = pm.getPackageInstaller();
+                PackageInstaller.SessionParams parentParams = new PackageInstaller.SessionParams(
                         PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-                // TODO: We can't get the installerPackageName for apex
-                // (b/123920130). Is it okay to ignore the installer package
-                // for apex?
-                if (!info.isApex()) {
-                    String installerPackageName = pm.getInstallerPackageName(info.getPackageName());
-                    if (installerPackageName != null) {
-                        params.setInstallerPackageName(installerPackageName);
-                    }
-                }
-                params.setRequestDowngrade(true);
-                params.setRequiredInstalledVersionCode(
-                        info.getVersionRolledBackFrom().getLongVersionCode());
+                parentParams.setRequestDowngrade(true);
+                parentParams.setMultiPackage();
                 if (rollback.isStaged()) {
-                    params.setStaged();
-                }
-                if (info.isApex()) {
-                    params.setInstallAsApex();
-                }
-                int sessionId = packageInstaller.createSession(params);
-                PackageInstaller.Session session = packageInstaller.openSession(sessionId);
-                File[] packageCodePaths = RollbackStore.getPackageCodePaths(
-                        rollback, info.getPackageName());
-                if (packageCodePaths == null) {
-                    sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE,
-                            "Backup copy of package inaccessible");
-                    return;
+                    parentParams.setStaged();
                 }
 
-                for (File packageCodePath : packageCodePaths) {
-                    try (ParcelFileDescriptor fd = ParcelFileDescriptor.open(packageCodePath,
+                int parentSessionId = packageInstaller.createSession(parentParams);
+                PackageInstaller.Session parentSession = packageInstaller.openSession(
+                        parentSessionId);
+
+                for (PackageRollbackInfo info : rollback.info.getPackages()) {
+                    PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
+                            PackageInstaller.SessionParams.MODE_FULL_INSTALL);
+                    // TODO: We can't get the installerPackageName for apex
+                    // (b/123920130). Is it okay to ignore the installer package
+                    // for apex?
+                    if (!info.isApex()) {
+                        String installerPackageName =
+                                pm.getInstallerPackageName(info.getPackageName());
+                        if (installerPackageName != null) {
+                            params.setInstallerPackageName(installerPackageName);
+                        }
+                    }
+                    params.setRequestDowngrade(true);
+                    params.setRequiredInstalledVersionCode(
+                            info.getVersionRolledBackFrom().getLongVersionCode());
+                    if (rollback.isStaged()) {
+                        params.setStaged();
+                    }
+                    if (info.isApex()) {
+                        params.setInstallAsApex();
+                    }
+                    int sessionId = packageInstaller.createSession(params);
+                    PackageInstaller.Session session = packageInstaller.openSession(sessionId);
+                    File[] packageCodePaths = RollbackStore.getPackageCodePaths(
+                            rollback, info.getPackageName());
+                    if (packageCodePaths == null) {
+                        sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE,
+                                "Backup copy of package inaccessible");
+                        return;
+                    }
+
+                    for (File packageCodePath : packageCodePaths) {
+                        try (ParcelFileDescriptor fd = ParcelFileDescriptor.open(packageCodePath,
                                 ParcelFileDescriptor.MODE_READ_ONLY)) {
-                        final long token = Binder.clearCallingIdentity();
-                        try {
-                            session.write(packageCodePath.getName(), 0, packageCodePath.length(),
-                                    fd);
-                        } finally {
-                            Binder.restoreCallingIdentity(token);
+                            final long token = Binder.clearCallingIdentity();
+                            try {
+                                session.write(packageCodePath.getName(), 0,
+                                        packageCodePath.length(),
+                                        fd);
+                            } finally {
+                                Binder.restoreCallingIdentity(token);
+                            }
                         }
                     }
+                    parentSession.addChildSessionId(sessionId);
                 }
-                parentSession.addChildSessionId(sessionId);
-            }
 
-            final LocalIntentReceiver receiver = new LocalIntentReceiver(
-                    (Intent result) -> {
-                        getHandler().post(() -> {
+                final LocalIntentReceiver receiver = new LocalIntentReceiver(
+                        (Intent result) -> getHandler().post(() -> {
 
                             int status = result.getIntExtra(PackageInstaller.EXTRA_STATUS,
                                     PackageInstaller.STATUS_FAILURE);
@@ -450,21 +465,22 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                                 // TODO: Should we just kill this rollback if
                                 // commit failed? Why would we expect commit
                                 // not to fail again?
-                                synchronized (mLock) {
-                                    // TODO: Could this cause a rollback to be
-                                    // resurrected if it should otherwise have
-                                    // expired by now?
+                                // TODO: Could this cause a rollback to be
+                                // resurrected if it should otherwise have
+                                // expired by now?
+                                synchronized (rollback.getLock()) {
                                     rollback.setAvailable();
                                     rollback.setRestoreUserDataInProgress(false);
                                 }
-                                sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE_INSTALL,
+                                sendFailure(statusReceiver,
+                                        RollbackManager.STATUS_FAILURE_INSTALL,
                                         "Rollback downgrade install failed: "
-                                        + result.getStringExtra(
+                                                + result.getStringExtra(
                                                 PackageInstaller.EXTRA_STATUS_MESSAGE));
                                 return;
                             }
 
-                            synchronized (mLock) {
+                            synchronized (rollback.getLock()) {
                                 if (!rollback.isStaged()) {
                                     // All calls to restoreUserData should have
                                     // completed by now for a non-staged install.
@@ -473,32 +489,31 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
                                 rollback.info.setCommittedSessionId(parentSessionId);
                                 rollback.info.getCausePackages().addAll(causePackages);
+                                RollbackStore.deletePackageCodePaths(rollback);
+                                saveRollback(rollback);
                             }
-                            mRollbackStore.deletePackageCodePaths(rollback);
-                            saveRollback(rollback);
 
                             sendSuccess(statusReceiver);
 
                             Intent broadcast = new Intent(Intent.ACTION_ROLLBACK_COMMITTED);
 
                             for (UserInfo userInfo : UserManager.get(mContext).getUsers(true)) {
-                                mContext.sendBroadcastAsUser(broadcast, userInfo.getUserHandle(),
+                                mContext.sendBroadcastAsUser(broadcast,
+                                        userInfo.getUserHandle(),
                                         Manifest.permission.MANAGE_ROLLBACKS);
                             }
-                        });
-                    }
-            );
+                        })
+                );
 
-            synchronized (mLock) {
                 rollback.setCommitted();
                 rollback.setRestoreUserDataInProgress(true);
+                parentSession.commit(receiver.getIntentSender());
+            } catch (IOException e) {
+                Slog.e(TAG, "Rollback failed", e);
+                sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE,
+                        "IOException: " + e.toString());
+                return;
             }
-            parentSession.commit(receiver.getIntentSender());
-        } catch (IOException e) {
-            Slog.e(TAG, "Rollback failed", e);
-            sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE,
-                    "IOException: " + e.toString());
-            return;
         }
     }
 
@@ -534,19 +549,23 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             Iterator<Rollback> iter = mRollbacks.iterator();
             while (iter.hasNext()) {
                 Rollback rollback = iter.next();
-                for (PackageRollbackInfo info : rollback.info.getPackages()) {
-                    if (info.getPackageName().equals(packageName)) {
-                        iter.remove();
-                        deleteRollback(rollback);
-                        break;
+                synchronized (rollback.getLock()) {
+                    for (PackageRollbackInfo info : rollback.info.getPackages()) {
+                        if (info.getPackageName().equals(packageName)) {
+                            iter.remove();
+                            deleteRollback(rollback);
+                            break;
+                        }
                     }
                 }
             }
             for (NewRollback newRollback : mNewRollbacks) {
-                for (PackageRollbackInfo info : newRollback.rollback.info.getPackages()) {
-                    if (info.getPackageName().equals(packageName)) {
-                        newRollback.isCancelled = true;
-                        break;
+                synchronized (newRollback.rollback.getLock()) {
+                    for (PackageRollbackInfo info : newRollback.rollback.info.getPackages()) {
+                        if (info.getPackageName().equals(packageName)) {
+                            newRollback.isCancelled = true;
+                            break;
+                        }
                     }
                 }
             }
@@ -580,8 +599,11 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
             for (int i = 0; i < rollbacks.size(); i++) {
                 Rollback rollback = rollbacks.get(i);
-                if (mAppDataRollbackHelper.commitPendingBackupAndRestoreForUser(userId, rollback)) {
-                    saveRollback(rollback);
+                synchronized (rollback.getLock()) {
+                    if (mAppDataRollbackHelper.commitPendingBackupAndRestoreForUser(
+                            userId, rollback)) {
+                        saveRollback(rollback);
+                    }
                 }
             }
 
@@ -618,16 +640,18 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             Set<String> apexPackageNames = new HashSet<>();
             synchronized (mLock) {
                 for (Rollback rollback : mRollbacks) {
-                    if (rollback.isStaged()) {
-                        if (rollback.isEnabling()) {
-                            enabling.add(rollback);
-                        } else if (rollback.isRestoreUserDataInProgress()) {
-                            restoreInProgress.add(rollback);
-                        }
+                    synchronized (rollback.getLock()) {
+                        if (rollback.isStaged()) {
+                            if (rollback.isEnabling()) {
+                                enabling.add(rollback);
+                            } else if (rollback.isRestoreUserDataInProgress()) {
+                                restoreInProgress.add(rollback);
+                            }
 
-                        for (PackageRollbackInfo info : rollback.info.getPackages()) {
-                            if (info.isApex()) {
-                                apexPackageNames.add(info.getPackageName());
+                            for (PackageRollbackInfo info : rollback.info.getPackages()) {
+                                if (info.isApex()) {
+                                    apexPackageNames.add(info.getPackageName());
+                                }
                             }
                         }
                     }
@@ -636,30 +660,32 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
             for (Rollback rollback : enabling) {
                 PackageInstaller installer = mContext.getPackageManager().getPackageInstaller();
-                PackageInstaller.SessionInfo session =
-                        installer.getSessionInfo(rollback.getStagedSessionId());
-                if (session == null || session.isStagedSessionFailed()) {
-                    // TODO: Do we need to remove this from
-                    // mRollbacks, or is it okay to leave as
-                    // unavailable until the next reboot when it will go
-                    // away on its own?
-                    deleteRollback(rollback);
-                } else if (session.isStagedSessionApplied()) {
-                    makeRollbackAvailable(rollback);
+                synchronized (rollback.getLock()) {
+                    PackageInstaller.SessionInfo session =
+                            installer.getSessionInfo(rollback.getStagedSessionId());
+                    if (session == null || session.isStagedSessionFailed()) {
+                        // TODO: Do we need to remove this from
+                        // mRollbacks, or is it okay to leave as
+                        // unavailable until the next reboot when it will go
+                        // away on its own?
+                        deleteRollback(rollback);
+                    } else if (session.isStagedSessionApplied()) {
+                        makeRollbackAvailable(rollback);
+                    }
                 }
             }
 
             for (Rollback rollback : restoreInProgress) {
                 PackageInstaller installer = mContext.getPackageManager().getPackageInstaller();
-                PackageInstaller.SessionInfo session =
-                        installer.getSessionInfo(rollback.getStagedSessionId());
-                // TODO: What if session is null?
-                if (session != null) {
-                    if (session.isStagedSessionApplied() || session.isStagedSessionFailed()) {
-                        synchronized (mLock) {
+                synchronized (rollback.getLock()) {
+                    PackageInstaller.SessionInfo session =
+                            installer.getSessionInfo(rollback.getStagedSessionId());
+                    // TODO: What if session is null?
+                    if (session != null) {
+                        if (session.isStagedSessionApplied() || session.isStagedSessionFailed()) {
                             rollback.setRestoreUserDataInProgress(false);
+                            saveRollback(rollback);
                         }
-                        saveRollback(rollback);
                     }
                 }
             }
@@ -694,16 +720,18 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             Iterator<Rollback> iter = mRollbacks.iterator();
             while (iter.hasNext()) {
                 Rollback rollback = iter.next();
-                // TODO: Should we remove rollbacks in the ENABLING state here?
-                if (rollback.isEnabling() || rollback.isAvailable()) {
-                    for (PackageRollbackInfo info : rollback.info.getPackages()) {
-                        if (info.getPackageName().equals(packageName)
-                                && !packageVersionsEqual(
+                synchronized (rollback.getLock()) {
+                    // TODO: Should we remove rollbacks in the ENABLING state here?
+                    if (rollback.isEnabling() || rollback.isAvailable()) {
+                        for (PackageRollbackInfo info : rollback.info.getPackages()) {
+                            if (info.getPackageName().equals(packageName)
+                                    && !packageVersionsEqual(
                                     info.getVersionRolledBackFrom(),
                                     installedVersion)) {
-                            iter.remove();
-                            deleteRollback(rollback);
-                            break;
+                                iter.remove();
+                                deleteRollback(rollback);
+                                break;
+                            }
                         }
                     }
                 }
@@ -761,16 +789,18 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             Iterator<Rollback> iter = mRollbacks.iterator();
             while (iter.hasNext()) {
                 Rollback rollback = iter.next();
-                if (!rollback.isAvailable()) {
-                    continue;
-                }
-                if (!now.isBefore(
+                synchronized (rollback.getLock()) {
+                    if (!rollback.isAvailable()) {
+                        continue;
+                    }
+                    if (!now.isBefore(
                             rollback.getTimestamp()
                                     .plusMillis(mRollbackLifetimeDurationInMillis))) {
-                    iter.remove();
-                    deleteRollback(rollback);
-                } else if (oldest == null || oldest.isAfter(rollback.getTimestamp())) {
-                    oldest = rollback.getTimestamp();
+                        iter.remove();
+                        deleteRollback(rollback);
+                    } else if (oldest == null || oldest.isAfter(rollback.getTimestamp())) {
+                        oldest = rollback.getTimestamp();
+                    }
                 }
             }
         }
@@ -878,10 +908,12 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         synchronized (mLock) {
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.getApkSessionId() == parentSession.getSessionId()) {
-                    // This is the apk session for a staged session with rollback enabled. We do not
-                    // need to create a new rollback for this session.
-                    return true;
+                synchronized (rollback.getLock()) {
+                    if (rollback.getApkSessionId() == parentSession.getSessionId()) {
+                        // This is the apk session for a staged session with rollback enabled. We do
+                        // not need to create a new rollback for this session.
+                        return true;
+                    }
                 }
             }
         }
@@ -980,6 +1012,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 new IntArray() /* pendingBackups */, new ArrayList<>() /* pendingRestores */,
                 isApex, new IntArray(), new SparseLongArray() /* ceSnapshotInodes */);
 
+
         try {
             ApplicationInfo appInfo = pkgInfo.applicationInfo;
             RollbackStore.backupPackageCodePath(rollback, packageName, appInfo.sourceDir);
@@ -993,7 +1026,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             return false;
         }
 
-        synchronized (mLock) {
+        synchronized (rollback.getLock()) {
             rollback.info.getPackages().add(packageRollbackInfo);
         }
         return true;
@@ -1021,27 +1054,31 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             // staged installs
             for (int i = 0; i < mRollbacks.size(); i++) {
                 Rollback rollback = mRollbacks.get(i);
-                if (!rollback.isEnabling()) {
-                    continue;
-                }
+                synchronized (rollback.getLock()) {
+                    if (!rollback.isEnabling()) {
+                        continue;
+                    }
 
-                for (PackageRollbackInfo info : rollback.info.getPackages()) {
-                    if (info.getPackageName().equals(packageName)) {
-                        mAppDataRollbackHelper.snapshotAppData(
-                                rollback.info.getRollbackId(), info, userIds);
-                        saveRollback(rollback);
-                        break;
+                    for (PackageRollbackInfo info : rollback.info.getPackages()) {
+                        if (info.getPackageName().equals(packageName)) {
+                            mAppDataRollbackHelper.snapshotAppData(
+                                    rollback.info.getRollbackId(), info, userIds);
+                            saveRollback(rollback);
+                            break;
+                        }
                     }
                 }
             }
             // non-staged installs
             PackageRollbackInfo info;
             for (NewRollback rollback : mNewRollbacks) {
-                info = getPackageRollbackInfo(rollback.rollback, packageName);
-                if (info != null) {
-                    mAppDataRollbackHelper.snapshotAppData(
-                            rollback.rollback.info.getRollbackId(), info, userIds);
-                    saveRollback(rollback.rollback);
+                synchronized (rollback.rollback.getLock()) {
+                    info = getPackageRollbackInfo(rollback.rollback, packageName);
+                    if (info != null) {
+                        mAppDataRollbackHelper.snapshotAppData(
+                                rollback.rollback.info.getRollbackId(), info, userIds);
+                        saveRollback(rollback.rollback);
+                    }
                 }
             }
         }
@@ -1054,11 +1091,13 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         synchronized (mLock) {
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback candidate = mRollbacks.get(i);
-                if (candidate.isRestoreUserDataInProgress()) {
-                    info = getPackageRollbackInfo(candidate, packageName);
-                    if (info != null) {
-                        rollback = candidate;
-                        break;
+                synchronized (candidate.getLock()) {
+                    if (candidate.isRestoreUserDataInProgress()) {
+                        info = getPackageRollbackInfo(candidate, packageName);
+                        if (info != null) {
+                            rollback = candidate;
+                            break;
+                        }
                     }
                 }
             }
@@ -1069,12 +1108,14 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         }
 
         for (int userId : userIds) {
-            final boolean changedRollback = mAppDataRollbackHelper.restoreAppData(
-                    rollback.info.getRollbackId(), info, userId, appId, seInfo);
+            synchronized (rollback.getLock()) {
+                final boolean changedRollback = mAppDataRollbackHelper.restoreAppData(
+                        rollback.info.getRollbackId(), info, userId, appId, seInfo);
 
-            // We've updated metadata about this rollback, so save it to flash.
-            if (changedRollback) {
-                saveRollback(rollback);
+                // We've updated metadata about this rollback, so save it to flash.
+                if (changedRollback) {
+                    saveRollback(rollback);
+                }
             }
         }
     }
@@ -1148,7 +1189,6 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
                 for (int i = 0; i < mRollbacks.size(); ++i) {
                     Rollback candidate = mRollbacks.get(i);
                     if (candidate.getStagedSessionId() == originalSessionId) {
-                        candidate.setApkSessionId(apkSessionId);
                         rollback = candidate;
                         break;
                     }
@@ -1163,7 +1203,10 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
             }
 
             if (rollback != null) {
-                saveRollback(rollback);
+                synchronized (rollback.getLock()) {
+                    rollback.setApkSessionId(apkSessionId);
+                    saveRollback(rollback);
+                }
             }
         });
     }
@@ -1274,44 +1317,49 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
 
             if (newRollback != null) {
                 Rollback rollback = completeEnableRollback(newRollback, success);
-                if (rollback != null && !rollback.isStaged()) {
-                    makeRollbackAvailable(rollback);
+                if (rollback != null) {
+                    synchronized (rollback.getLock()) {
+                        if (!rollback.isStaged()) {
+                            makeRollbackAvailable(rollback);
+                        }
+                    }
                 }
             }
         }
     }
 
     /**
-     * Add a rollback to the list of rollbacks.
-     * This should be called after rollback has been enabled for all packages
-     * in the rollback. It does not make the rollback available yet.
+     * Add a rollback to the list of rollbacks. This should be called after rollback has been
+     * enabled for all packages in the rollback. It does not make the rollback available yet.
+     *
+     * <p>Note that no rollback-specific locks should be held when this method is called.
      *
      * @return the Rollback instance for a successfully enable-completed rollback,
      * or null on error.
      */
     private Rollback completeEnableRollback(NewRollback newRollback, boolean success) {
         Rollback rollback = newRollback.rollback;
-        if (!success) {
-            // The install session was aborted, clean up the pending install.
-            deleteRollback(rollback);
-            return null;
-        }
-        if (newRollback.isCancelled) {
-            Slog.e(TAG, "Rollback has been cancelled by PackageManager");
-            deleteRollback(rollback);
-            return null;
-        }
+        synchronized (rollback.getLock()) {
+            if (!success) {
+                // The install session was aborted, clean up the pending install.
+                deleteRollback(rollback);
+                return null;
+            }
+            if (newRollback.isCancelled) {
+                Slog.e(TAG, "Rollback has been cancelled by PackageManager");
+                deleteRollback(rollback);
+                return null;
+            }
 
-        // It's safe to access rollback.info outside a synchronized block because
-        // this is running on the handler thread and all changes to the
-        // rollback.info occur on the handler thread.
-        if (rollback.info.getPackages().size() != newRollback.packageSessionIds.length) {
-            Slog.e(TAG, "Failed to enable rollback for all packages in session.");
-            deleteRollback(rollback);
-            return null;
-        }
 
-        saveRollback(rollback);
+            if (rollback.info.getPackages().size() != newRollback.packageSessionIds.length) {
+                Slog.e(TAG, "Failed to enable rollback for all packages in session.");
+                deleteRollback(rollback);
+                return null;
+            }
+
+            saveRollback(rollback);
+        }
         synchronized (mLock) {
             // Note: There is a small window of time between when
             // the session has been committed by the package
@@ -1329,14 +1377,13 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         return rollback;
     }
 
+    @GuardedBy("rollback.getLock")
     private void makeRollbackAvailable(Rollback rollback) {
         // TODO: What if the rollback has since been expired, for example due
         // to a new package being installed. Won't this revive an expired
         // rollback? Consider adding a ROLLBACK_STATE_EXPIRED to address this.
-        synchronized (mLock) {
-            rollback.setAvailable();
-            rollback.setTimestamp(Instant.now());
-        }
+        rollback.setAvailable();
+        rollback.setTimestamp(Instant.now());
         saveRollback(rollback);
 
         // TODO(zezeozue): Provide API to explicitly start observing instead
@@ -1373,6 +1420,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
      * Returns the {@code PackageRollbackInfo} associated with {@code packageName} from
      * a specified {@code Rollback}.
      */
+    @GuardedBy("rollback.getLock")
     private static PackageRollbackInfo getPackageRollbackInfo(Rollback rollback,
             String packageName) {
         for (PackageRollbackInfo info : rollback.info.getPackages()) {
@@ -1399,6 +1447,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         throw new IllegalStateException("Failed to allocate rollback ID");
     }
 
+    @GuardedBy("rollback.getLock")
     private void deleteRollback(Rollback rollback) {
         for (PackageRollbackInfo info : rollback.info.getPackages()) {
             IntArray snapshottedUsers = info.getSnapshottedUsers();
@@ -1417,6 +1466,7 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
      * TODO: Double check we can't do a better job handling the IOException in
      * a cases where this method is called.
      */
+    @GuardedBy("rollback.getLock")
     private void saveRollback(Rollback rollback) {
         try {
             mRollbackStore.saveRollback(rollback);
@@ -1431,32 +1481,34 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         IndentingPrintWriter ipw = new IndentingPrintWriter(pw, "  ");
         synchronized (mLock) {
             for (Rollback rollback : mRollbacks) {
-                RollbackInfo info = rollback.info;
-                ipw.println(info.getRollbackId() + ":");
-                ipw.increaseIndent();
-                ipw.println("-state: " + rollback.getStateAsString());
-                ipw.println("-timestamp: " + rollback.getTimestamp());
-                if (rollback.getStagedSessionId() != -1) {
-                    ipw.println("-stagedSessionId: " + rollback.getStagedSessionId());
-                }
-                ipw.println("-packages:");
-                ipw.increaseIndent();
-                for (PackageRollbackInfo pkg : info.getPackages()) {
-                    ipw.println(pkg.getPackageName()
-                            + " " + pkg.getVersionRolledBackFrom().getLongVersionCode()
-                            + " -> " + pkg.getVersionRolledBackTo().getLongVersionCode());
-                }
-                ipw.decreaseIndent();
-                if (rollback.isCommitted()) {
-                    ipw.println("-causePackages:");
+                synchronized (rollback.getLock()) {
+                    RollbackInfo info = rollback.info;
+                    ipw.println(info.getRollbackId() + ":");
+                    ipw.increaseIndent();
+                    ipw.println("-state: " + rollback.getStateAsString());
+                    ipw.println("-timestamp: " + rollback.getTimestamp());
+                    if (rollback.getStagedSessionId() != -1) {
+                        ipw.println("-stagedSessionId: " + rollback.getStagedSessionId());
+                    }
+                    ipw.println("-packages:");
                     ipw.increaseIndent();
-                    for (VersionedPackage cPkg : info.getCausePackages()) {
-                        ipw.println(cPkg.getPackageName() + " " + cPkg.getLongVersionCode());
+                    for (PackageRollbackInfo pkg : info.getPackages()) {
+                        ipw.println(pkg.getPackageName()
+                                + " " + pkg.getVersionRolledBackFrom().getLongVersionCode()
+                                + " -> " + pkg.getVersionRolledBackTo().getLongVersionCode());
+                    }
+                    ipw.decreaseIndent();
+                    if (rollback.isCommitted()) {
+                        ipw.println("-causePackages:");
+                        ipw.increaseIndent();
+                        for (VersionedPackage cPkg : info.getCausePackages()) {
+                            ipw.println(cPkg.getPackageName() + " " + cPkg.getLongVersionCode());
+                        }
+                        ipw.decreaseIndent();
+                        ipw.println("-committedSessionId: " + info.getCommittedSessionId());
                     }
                     ipw.decreaseIndent();
-                    ipw.println("-committedSessionId: " + info.getCommittedSessionId());
                 }
-                ipw.decreaseIndent();
             }
         }
     }
@@ -1517,7 +1569,8 @@ class RollbackManagerServiceImpl extends IRollbackManager.Stub {
         }
     }
 
-    NewRollback createNewRollbackLocked(PackageInstaller.SessionInfo parentSession) {
+    @GuardedBy("mLock")
+    private NewRollback createNewRollbackLocked(PackageInstaller.SessionInfo parentSession) {
         int rollbackId = allocateRollbackIdLocked();
         final Rollback rollback;
         int parentSessionId = parentSession.getSessionId();
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index 772c53fec4ce..b6d1f1875907 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -27,6 +27,8 @@ import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseLongArray;
 
+import com.android.internal.annotations.GuardedBy;
+
 import libcore.io.IoUtils;
 
 import org.json.JSONArray;
@@ -250,6 +252,7 @@ class RollbackStore {
     /**
      * Saves the given rollback to persistent storage.
      */
+    @GuardedBy("rollback.getLock")
     void saveRollback(Rollback rollback) throws IOException {
         try {
             JSONObject dataJson = new JSONObject();