Merge "Simplify staged install resume flow." into sc-dev

5 files changed, 964 insertions, 298 deletions

diff --git a/services/core/java/com/android/server/pm/ApexManager.java b/services/core/java/com/android/server/pm/ApexManager.java
index de85d9e25642..f31d1da84c35 100644
--- a/services/core/java/com/android/server/pm/ApexManager.java
+++ b/services/core/java/com/android/server/pm/ApexManager.java
@@ -43,6 +43,7 @@ import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Singleton;
 import android.util.Slog;
+import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -226,6 +227,12 @@ public abstract class ApexManager {
     abstract ApexSessionInfo getStagedSessionInfo(int sessionId);
 
     /**
+     * Returns array of all staged sessions known to apexd.
+     */
+    @NonNull
+    abstract SparseArray<ApexSessionInfo> getSessions();
+
+    /**
      * Submit a staged session to apex service. This causes the apex service to perform some initial
      * verification and accept or reject the session. Submitting a session successfully is not
      * enough for it to be activated at the next boot, the caller needs to call
@@ -691,6 +698,21 @@ public abstract class ApexManager {
         }
 
         @Override
+        SparseArray<ApexSessionInfo> getSessions() {
+            try {
+                final ApexSessionInfo[] sessions = waitForApexService().getSessions();
+                final SparseArray<ApexSessionInfo> result = new SparseArray<>(sessions.length);
+                for (int i = 0; i < sessions.length; i++) {
+                    result.put(sessions[i].sessionId, sessions[i]);
+                }
+                return result;
+            } catch (RemoteException re) {
+                Slog.e(TAG, "Unable to contact apexservice", re);
+                throw new RuntimeException(re);
+            }
+        }
+
+        @Override
         ApexInfoList submitStagedSession(ApexSessionParams params) throws PackageManagerException {
             try {
                 final ApexInfoList apexInfoList = new ApexInfoList();
@@ -1083,6 +1105,11 @@ public abstract class ApexManager {
         }
 
         @Override
+        SparseArray<ApexSessionInfo> getSessions() {
+            return new SparseArray<>(0);
+        }
+
+        @Override
         ApexInfoList submitStagedSession(ApexSessionParams params)
                 throws PackageManagerException {
             throw new PackageManagerException(PackageManager.INSTALL_FAILED_INTERNAL_ERROR,
diff --git a/services/core/java/com/android/server/pm/PackageInstallerService.java b/services/core/java/com/android/server/pm/PackageInstallerService.java
index 2d393c089411..281283048a95 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerService.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerService.java
@@ -295,23 +295,29 @@ public class PackageInstallerService extends IPackageInstaller.Stub implements
         synchronized (mSessions) {
             for (int i = 0; i < mSessions.size(); i++) {
                 final PackageInstallerSession session = mSessions.valueAt(i);
-                if (session.isStaged()) {
-                    stagedSessionsToRestore.add(session.mStagedSession);
+                if (!session.isStaged()) {
+                    continue;
+                }
+                StagingManager.StagedSession stagedSession = session.mStagedSession;
+                if (!stagedSession.isInTerminalState() && stagedSession.hasParentSessionId()
+                        && getSession(stagedSession.getParentSessionId()) == null) {
+                    stagedSession.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
+                            "An orphan staged session " + stagedSession.sessionId() + " is found, "
+                                + "parent " + stagedSession.getParentSessionId() + " is missing");
+                    continue;
+                }
+                if (!stagedSession.hasParentSessionId() && stagedSession.isCommitted()
+                        && !stagedSession.isInTerminalState()) {
+                    // StagingManager.restoreSessions expects a list of committed, non-finalized
+                    // parent staged sessions.
+                    stagedSessionsToRestore.add(stagedSession);
                 }
             }
         }
-        // Don't hold mSessions lock when calling restoreSession, since it might trigger an APK
+        // Don't hold mSessions lock when calling restoreSessions, since it might trigger an APK
         // atomic install which needs to query sessions, which requires lock on mSessions.
-        boolean isDeviceUpgrading = mPm.isDeviceUpgrading();
-        for (StagingManager.StagedSession session : stagedSessionsToRestore) {
-            if (!session.isInTerminalState() && session.hasParentSessionId()
-                    && getSession(session.getParentSessionId()) == null) {
-                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
-                        "An orphan staged session " + session.sessionId() + " is found, "
-                                + "parent " + session.getParentSessionId() + " is missing");
-            }
-            mStagingManager.restoreSession(session, isDeviceUpgrading);
-        }
+        // Note: restoreSessions mutates content of stagedSessionsToRestore.
+        mStagingManager.restoreSessions(stagedSessionsToRestore, mPm.isDeviceUpgrading());
     }
 
     @GuardedBy("mSessions")
diff --git a/services/core/java/com/android/server/pm/StagingManager.java b/services/core/java/com/android/server/pm/StagingManager.java
index 545567c26972..0a74032ab214 100644
--- a/services/core/java/com/android/server/pm/StagingManager.java
+++ b/services/core/java/com/android/server/pm/StagingManager.java
@@ -50,8 +50,6 @@ import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.SystemProperties;
 import android.os.UserHandle;
-import android.os.storage.IStorageManager;
-import android.os.storage.StorageManager;
 import android.text.TextUtils;
 import android.util.ArraySet;
 import android.util.IntArray;
@@ -63,6 +61,7 @@ import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.content.PackageHelper;
 import com.android.internal.os.BackgroundThread;
+import com.android.internal.util.Preconditions;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
 import com.android.server.SystemServiceManager;
@@ -143,10 +142,16 @@ public class StagingManager {
     }
 
     StagingManager(Context context, Supplier<PackageParser2> packageParserSupplier) {
+        this(context, packageParserSupplier, ApexManager.getInstance());
+    }
+
+    @VisibleForTesting
+    StagingManager(Context context, Supplier<PackageParser2> packageParserSupplier,
+            ApexManager apexManager) {
         mContext = context;
         mPackageParserSupplier = packageParserSupplier;
 
-        mApexManager = ApexManager.getInstance();
+        mApexManager = apexManager;
         mPowerManager = (PowerManager) context.getSystemService(Context.POWER_SERVICE);
         mPreRebootVerificationHandler = new PreRebootVerificationHandler(
                 BackgroundThread.get().getLooper());
@@ -354,11 +359,11 @@ public class StagingManager {
     }
 
     // Reverts apex sessions and user data (if checkpoint is supported). Also reboots the device.
-    private void abortCheckpoint(int sessionId, String errorMsg) {
-        String failureReason = "Failed to install sessionId: " + sessionId + " Error: " + errorMsg;
+    private void abortCheckpoint(String failureReason, boolean supportsCheckpoint,
+            boolean needsCheckpoint) {
         Slog.e(TAG, failureReason);
         try {
-            if (supportsCheckpoint() && needsCheckpoint()) {
+            if (supportsCheckpoint && needsCheckpoint) {
                 // Store failure reason for next reboot
                 try (BufferedWriter writer =
                              new BufferedWriter(new FileWriter(mFailureReasonFile))) {
@@ -371,8 +376,9 @@ public class StagingManager {
                 if (mApexManager.isApexSupported()) {
                     mApexManager.revertActiveSessions();
                 }
+
                 PackageHelper.getStorageManager().abortChanges(
-                        "StagingManager initiated", false /*retry*/);
+                        "abort-staged-install", false /*retry*/);
             }
         } catch (Exception e) {
             Slog.wtf(TAG, "Failed to abort checkpoint", e);
@@ -384,14 +390,6 @@ public class StagingManager {
         }
     }
 
-    private boolean supportsCheckpoint() throws RemoteException {
-        return PackageHelper.getStorageManager().supportsCheckpoint();
-    }
-
-    private boolean needsCheckpoint() throws RemoteException {
-        return PackageHelper.getStorageManager().needsCheckpoint();
-    }
-
     /**
      * Utility function for extracting apex sessions out of multi-package/single session.
      */
@@ -517,96 +515,31 @@ public class StagingManager {
         }
     }
 
-    private void resumeSession(@NonNull StagedSession session)
-            throws PackageManagerException {
+    private void resumeSession(@NonNull StagedSession session, boolean supportsCheckpoint,
+            boolean needsCheckpoint) throws PackageManagerException {
         Slog.d(TAG, "Resuming session " + session.sessionId());
 
         final boolean hasApex = session.containsApexSession();
-        ApexSessionInfo apexSessionInfo = null;
-        if (hasApex) {
-            // Check with apexservice whether the apex packages have been activated.
-            apexSessionInfo = mApexManager.getStagedSessionInfo(session.sessionId());
-
-            // Prepare for logging a native crash during boot, if one occurred.
-            if (apexSessionInfo != null && !TextUtils.isEmpty(
-                    apexSessionInfo.crashingNativeProcess)) {
-                prepareForLoggingApexdRevert(session, apexSessionInfo.crashingNativeProcess);
-            }
-
-            if (apexSessionInfo != null && apexSessionInfo.isVerified) {
-                // Session has been previously submitted to apexd, but didn't complete all the
-                // pre-reboot verification, perhaps because the device rebooted in the meantime.
-                // Greedily re-trigger the pre-reboot verification. We want to avoid marking it as
-                // failed when not in checkpoint mode, hence it is being processed separately.
-                Slog.d(TAG, "Found pending staged session " + session.sessionId() + " still to "
-                        + "be verified, resuming pre-reboot verification");
-                mPreRebootVerificationHandler.startPreRebootVerification(session);
-                return;
-            }
-        }
 
         // Before we resume session, we check if revert is needed or not. Typically, we enter file-
         // system checkpoint mode when we reboot first time in order to install staged sessions. We
         // want to install staged sessions in this mode as rebooting now will revert user data. If
         // something goes wrong, then we reboot again to enter fs-rollback mode. Rebooting now will
         // have no effect on user data, so mark the sessions as failed instead.
-        try {
-            // If checkpoint is supported, then we only resume sessions if we are in checkpointing
-            // mode. If not, we fail all sessions.
-            if (supportsCheckpoint() && !needsCheckpoint()) {
-                String revertMsg = "Reverting back to safe state. Marking "
-                        + session.sessionId() + " as failed.";
-                final String reasonForRevert = getReasonForRevert();
-                if (!TextUtils.isEmpty(reasonForRevert)) {
-                    revertMsg += " Reason for revert: " + reasonForRevert;
-                }
-                Slog.d(TAG, revertMsg);
-                session.setSessionFailed(SessionInfo.STAGED_SESSION_UNKNOWN, revertMsg);
-                return;
+        // If checkpoint is supported, then we only resume sessions if we are in checkpointing mode.
+        // If not, we fail all sessions.
+        if (supportsCheckpoint && !needsCheckpoint) {
+            String revertMsg = "Reverting back to safe state. Marking " + session.sessionId()
+                    + " as failed.";
+            final String reasonForRevert = getReasonForRevert();
+            if (!TextUtils.isEmpty(reasonForRevert)) {
+                revertMsg += " Reason for revert: " + reasonForRevert;
             }
-        } catch (RemoteException e) {
-            // Cannot continue staged install without knowing if fs-checkpoint is supported
-            Slog.e(TAG, "Checkpoint support unknown. Aborting staged install for session "
-                    + session.sessionId(), e);
-            // TODO: Mark all staged sessions together and reboot only once
-            session.setSessionFailed(SessionInfo.STAGED_SESSION_UNKNOWN,
-                    "Checkpoint support unknown. Aborting staged install.");
-            if (hasApex) {
-                mApexManager.revertActiveSessions();
-            }
-            mPowerManager.reboot("Checkpoint support unknown");
+            Slog.d(TAG, revertMsg);
+            session.setSessionFailed(SessionInfo.STAGED_SESSION_UNKNOWN, revertMsg);
             return;
         }
 
-        // Check if apex packages in the session failed to activate
-        if (hasApex) {
-            if (apexSessionInfo == null) {
-                final String errorMsg = "apexd did not know anything about a staged session "
-                        + "supposed to be activated";
-                throw new PackageManagerException(
-                        SessionInfo.STAGED_SESSION_ACTIVATION_FAILED, errorMsg);
-            }
-            if (isApexSessionFailed(apexSessionInfo)) {
-                String errorMsg = "APEX activation failed. Check logcat messages from apexd "
-                        + "for more information.";
-                if (!TextUtils.isEmpty(mNativeFailureReason)) {
-                    errorMsg = "Session reverted due to crashing native process: "
-                            + mNativeFailureReason;
-                }
-                throw new PackageManagerException(
-                        SessionInfo.STAGED_SESSION_ACTIVATION_FAILED, errorMsg);
-            }
-            if (!apexSessionInfo.isActivated && !apexSessionInfo.isSuccess) {
-                // Apexd did not apply the session for some unknown reason. There is no
-                // guarantee that apexd will install it next time. Safer to proactively mark
-                // it as failed.
-                final String errorMsg = "Staged session " + session.sessionId() + "at boot "
-                        + "didn't activate nor fail. Marking it as failed anyway.";
-                throw new PackageManagerException(
-                        SessionInfo.STAGED_SESSION_ACTIVATION_FAILED, errorMsg);
-            }
-        }
-
         // Handle apk and apk-in-apex installation
         if (hasApex) {
             checkInstallationOfApkInApexSuccessful(session);
@@ -622,28 +555,24 @@ public class StagingManager {
         Slog.d(TAG, "Marking session " + session.sessionId() + " as applied");
         session.setSessionApplied();
         if (hasApex) {
-            try {
-                if (supportsCheckpoint()) {
-                    // Store the session ID, which will be marked as successful by ApexManager
-                    // upon boot completion.
-                    synchronized (mSuccessfulStagedSessionIds) {
-                        mSuccessfulStagedSessionIds.add(session.sessionId());
-                    }
-                } else {
-                    // Mark sessions as successful immediately on non-checkpointing devices.
-                    mApexManager.markStagedSessionSuccessful(session.sessionId());
+            if (supportsCheckpoint) {
+                // Store the session ID, which will be marked as successful by ApexManager upon
+                // boot completion.
+                synchronized (mSuccessfulStagedSessionIds) {
+                    mSuccessfulStagedSessionIds.add(session.sessionId());
                 }
-            } catch (RemoteException e) {
-                Slog.w(TAG, "Checkpoint support unknown, marking session as successful "
-                        + "immediately.");
+            } else {
+                // Mark sessions as successful immediately on non-checkpointing devices.
                 mApexManager.markStagedSessionSuccessful(session.sessionId());
             }
         }
     }
 
-    void onInstallationFailure(StagedSession session, PackageManagerException e) {
+    void onInstallationFailure(StagedSession session, PackageManagerException e,
+            boolean supportsCheckpoint, boolean needsCheckpoint) {
         session.setSessionFailed(e.error, e.getMessage());
-        abortCheckpoint(session.sessionId(), e.getMessage());
+        abortCheckpoint("Failed to install sessionId: " + session.sessionId()
+                + " Error: " + e.getMessage(), supportsCheckpoint, needsCheckpoint);
 
         // If checkpoint is not supported, we have to handle failure for one staged session.
         if (!session.containsApexSession()) {
@@ -767,8 +696,13 @@ public class StagingManager {
                     "Cannot stage session " + session.sessionId() + " with package name null");
         }
 
-        boolean supportsCheckpoint = ((StorageManager) mContext.getSystemService(
-                Context.STORAGE_SERVICE)).isCheckpointSupported();
+        boolean supportsCheckpoint;
+        try {
+            supportsCheckpoint = PackageHelper.getStorageManager().supportsCheckpoint();
+        } catch (RemoteException e) {
+            throw new PackageManagerException(SessionInfo.STAGED_SESSION_VERIFICATION_FAILED,
+                    "Can't query fs-checkpoint status : " + e);
+        }
 
         final boolean isRollback = isRollback(session);
 
@@ -911,60 +845,166 @@ public class StagingManager {
                 || apexSessionInfo.isRevertFailed;
     }
 
-    void restoreSession(@NonNull StagedSession session, boolean isDeviceUpgrading) {
-        if (session.hasParentSessionId()) {
-            // Only parent sessions can be restored
-            return;
-        }
-        // Store this parent session which will be used to check overlapping later
-        createSession(session);
-        // The preconditions used during pre-reboot verification might have changed when device
-        // is upgrading. Updated staged sessions to activation failed before we resume the session.
-        StagedSession sessionToResume = session;
-        if (isDeviceUpgrading && !sessionToResume.isInTerminalState()) {
-            sessionToResume.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
-                        "Build fingerprint has changed");
-            return;
+    private void handleNonReadyAndDestroyedSessions(List<StagedSession> sessions) {
+        int j = sessions.size();
+        for (int i = 0; i < j; ) {
+            // Maintain following invariant:
+            //  * elements at positions [0, i) should be kept
+            //  * elements at positions [j, n) should be remove.
+            //  * n = sessions.size()
+            StagedSession session = sessions.get(i);
+            if (session.isDestroyed()) {
+                // Device rebooted before abandoned session was cleaned up.
+                session.abandon();
+                StagedSession session2 = sessions.set(j - 1, session);
+                sessions.set(i, session2);
+                j--;
+            } else if (!session.isSessionReady()) {
+                // The framework got restarted before the pre-reboot verification could complete,
+                // restart the verification.
+                mPreRebootVerificationHandler.startPreRebootVerification(session);
+                StagedSession session2 = sessions.set(j - 1, session);
+                sessions.set(i, session2);
+                j--;
+            } else {
+                i++;
+            }
         }
-        checkStateAndResume(sessionToResume);
+        // Delete last j elements.
+        sessions.subList(j, sessions.size()).clear();
     }
 
-    private void checkStateAndResume(@NonNull StagedSession session) {
-        // Do not resume session if boot completed already
+    void restoreSessions(@NonNull List<StagedSession> sessions, boolean isDeviceUpgrading) {
+        // Do not resume sessions if boot completed already
         if (SystemProperties.getBoolean("sys.boot_completed", false)) {
             return;
         }
 
-        if (!session.isCommitted()) {
-            // Session hasn't been committed yet, ignore.
+        for (int i = 0; i < sessions.size(); i++) {
+            StagedSession session = sessions.get(i);
+            // Quick check that PackageInstallerService gave us sessions we expected.
+            Preconditions.checkArgument(!session.hasParentSessionId(),
+                    session.sessionId() + " is a child session");
+            Preconditions.checkArgument(session.isCommitted(),
+                    session.sessionId() + " is not committed");
+            Preconditions.checkArgument(!session.isInTerminalState(),
+                    session.sessionId() + " is in terminal state");
+            // Store this parent session which will be used to check overlapping later
+            createSession(session);
+        }
+
+        if (isDeviceUpgrading) {
+            // TODO(ioffe): check that corresponding apex sessions are failed.
+            // The preconditions used during pre-reboot verification might have changed when device
+            // is upgrading. Fail all the sessions and exit early.
+            for (int i = 0; i < sessions.size(); i++) {
+                StagedSession session = sessions.get(i);
+                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
+                        "Build fingerprint has changed");
+            }
             return;
         }
-        // Check the state of the session and decide what to do next.
-        if (session.isSessionFailed() || session.isSessionApplied()) {
-            // Final states, nothing to do.
+
+        boolean needsCheckpoint = false;
+        boolean supportsCheckpoint = false;
+        try {
+            supportsCheckpoint = PackageHelper.getStorageManager().supportsCheckpoint();
+            needsCheckpoint = PackageHelper.getStorageManager().needsCheckpoint();
+        } catch (RemoteException e) {
+            // This means that vold has crashed, and device is in a bad state.
+            throw new IllegalStateException("Failed to get checkpoint status", e);
+        }
+
+        if (sessions.size() > 1 && !supportsCheckpoint) {
+            throw new IllegalStateException("Detected multiple staged sessions on a device without "
+                    + "fs-checkpoint support");
+        }
+
+        // Do a set of quick checks before resuming individual sessions:
+        //   1. Schedule a pre-reboot verification for non-ready sessions.
+        //   2. Abandon destroyed sessions.
+        handleNonReadyAndDestroyedSessions(sessions); // mutates |sessions|
+
+        //   3. Check state of apex sessions is consistent. All non-applied sessions will be marked
+        //      as failed.
+        final SparseArray<ApexSessionInfo> apexSessions = mApexManager.getSessions();
+        boolean hasFailedApexSession = false;
+        boolean hasAppliedApexSession = false;
+        for (int i = 0; i < sessions.size(); i++) {
+            StagedSession session = sessions.get(i);
+            if (!session.containsApexSession()) {
+                // At this point we are only interested in apex sessions.
+                continue;
+            }
+            final ApexSessionInfo apexSession = apexSessions.get(session.sessionId());
+            if (apexSession == null || apexSession.isUnknown) {
+                hasFailedApexSession = true;
+                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED, "apexd did "
+                        + "not know anything about a staged session supposed to be activated");
+                continue;
+            } else if (isApexSessionFailed(apexSession)) {
+                hasFailedApexSession = true;
+                String errorMsg = "APEX activation failed. Check logcat messages from apexd "
+                        + "for more information.";
+                if (!TextUtils.isEmpty(apexSession.crashingNativeProcess)) {
+                    prepareForLoggingApexdRevert(session, apexSession.crashingNativeProcess);
+                    errorMsg = "Session reverted due to crashing native process: "
+                            + apexSession.crashingNativeProcess;
+                }
+                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED, errorMsg);
+                continue;
+            } else if (apexSession.isActivated || apexSession.isSuccess) {
+                hasAppliedApexSession = true;
+                continue;
+            } else if (apexSession.isStaged) {
+                // Apexd did not apply the session for some unknown reason. There is no guarantee
+                // that apexd will install it next time. Safer to proactively mark it as failed.
+                hasFailedApexSession = true;
+                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
+                        "Staged session " + session.sessionId() + " at boot didn't activate nor "
+                        + "fail. Marking it as failed anyway.");
+            } else {
+                Slog.w(TAG, "Apex session " + session.sessionId() + " is in impossible state");
+                hasFailedApexSession = true;
+                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
+                        "Impossible state");
+            }
+        }
+
+        if (hasAppliedApexSession && hasFailedApexSession) {
+            abortCheckpoint("Found both applied and failed apex sessions", supportsCheckpoint,
+                    needsCheckpoint);
             return;
         }
-        if (session.isDestroyed()) {
-            // Device rebooted before abandoned session was cleaned up.
-            session.abandon();
+
+        if (hasFailedApexSession) {
+            // Either of those means that we failed at least one apex session, hence we should fail
+            // all other sessions.
+            for (int i = 0; i < sessions.size(); i++) {
+                StagedSession session = sessions.get(i);
+                if (session.isSessionFailed()) {
+                    // Session has been already failed in the loop above.
+                    continue;
+                }
+                session.setSessionFailed(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
+                        "Another apex session failed");
+            }
             return;
         }
-        if (!session.isSessionReady()) {
-            // The framework got restarted before the pre-reboot verification could complete,
-            // restart the verification.
-            mPreRebootVerificationHandler.startPreRebootVerification(session);
-        } else {
-            // Session had already being marked ready. Start the checks to verify if there is any
-            // follow-up work.
+
+        // Time to resume sessions.
+        for (int i = 0; i < sessions.size(); i++) {
+            StagedSession session = sessions.get(i);
             try {
-                resumeSession(session);
+                resumeSession(session, supportsCheckpoint, needsCheckpoint);
             } catch (PackageManagerException e) {
-                onInstallationFailure(session, e);
+                onInstallationFailure(session, e, supportsCheckpoint, needsCheckpoint);
             } catch (Exception e) {
                 Slog.e(TAG, "Staged install failed due to unhandled exception", e);
                 onInstallationFailure(session, new PackageManagerException(
                         SessionInfo.STAGED_SESSION_ACTIVATION_FAILED,
-                        "Staged install failed due to unhandled exception: " + e));
+                        "Staged install failed due to unhandled exception: " + e),
+                        supportsCheckpoint, needsCheckpoint);
             }
         }
     }
@@ -992,9 +1032,7 @@ public class StagingManager {
         mContext.registerReceiver(new BroadcastReceiver() {
             @Override
             public void onReceive(Context ctx, Intent intent) {
-                mPreRebootVerificationHandler.readyToStart();
-                BackgroundThread.getExecutor().execute(
-                        () -> logFailedApexSessionsIfNecessary());
+                onBootCompletedBroadcastReceived();
                 ctx.unregisterReceiver(this);
             }
         }, new IntentFilter(Intent.ACTION_BOOT_COMPLETED));
@@ -1002,6 +1040,12 @@ public class StagingManager {
         mFailureReasonFile.delete();
     }
 
+    @VisibleForTesting
+    void onBootCompletedBroadcastReceived() {
+        mPreRebootVerificationHandler.readyToStart();
+        BackgroundThread.getExecutor().execute(() -> logFailedApexSessionsIfNecessary());
+    }
+
     private static class LocalIntentReceiverSync {
         private final LinkedBlockingQueue<Intent> mResult = new LinkedBlockingQueue<>();
 
@@ -1286,9 +1330,8 @@ public class StagingManager {
         private void handlePreRebootVerification_End(@NonNull StagedSession session) {
             // Before marking the session as ready, start checkpoint service if available
             try {
-                IStorageManager storageManager = PackageHelper.getStorageManager();
-                if (storageManager.supportsCheckpoint()) {
-                    storageManager.startCheckpoint(2);
+                if (PackageHelper.getStorageManager().supportsCheckpoint()) {
+                    PackageHelper.getStorageManager().startCheckpoint(2);
                 }
             } catch (Exception e) {
                 // Failed to get hold of StorageManager
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
new file mode 100644
index 000000000000..195cc010c9a7
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/StagingManagerTest.java
@@ -0,0 +1,725 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.pm;
+
+import android.apex.ApexSessionInfo;
+import android.content.Context;
+import android.content.IntentSender;
+import android.content.pm.PackageInstaller;
+import android.content.pm.PackageInstaller.SessionInfo;
+import android.content.pm.PackageInstaller.SessionInfo.StagedSessionErrorCode;
+import android.os.SystemProperties;
+import android.os.storage.IStorageManager;
+import android.platform.test.annotations.Presubmit;
+import android.util.SparseArray;
+
+import com.android.dx.mockito.inline.extended.ExtendedMockito;
+import com.android.internal.content.PackageHelper;
+import com.android.internal.os.BackgroundThread;
+import com.android.internal.util.Preconditions;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
+
+import java.io.File;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.function.Predicate;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertThrows;
+
+@Presubmit
+@RunWith(JUnit4.class)
+public class StagingManagerTest {
+    @Rule
+    public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
+
+    @Mock private Context mContext;
+    @Mock private IStorageManager mStorageManager;
+    @Mock private ApexManager mApexManager;
+
+    private File mTmpDir;
+    private StagingManager mStagingManager;
+
+    private MockitoSession mMockitoSession;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getSystemService(eq(Context.POWER_SERVICE))).thenReturn(null);
+
+        mMockitoSession = ExtendedMockito.mockitoSession()
+                    .strictness(Strictness.LENIENT)
+                    .mockStatic(SystemProperties.class)
+                    .mockStatic(PackageHelper.class)
+                    .startMocking();
+
+        when(mStorageManager.supportsCheckpoint()).thenReturn(true);
+        when(mStorageManager.needsCheckpoint()).thenReturn(true);
+        when(PackageHelper.getStorageManager()).thenReturn(mStorageManager);
+
+        when(SystemProperties.get(eq("ro.apex.updatable"))).thenReturn("true");
+        when(SystemProperties.get(eq("ro.apex.updatable"), anyString())).thenReturn("true");
+
+        mTmpDir = mTemporaryFolder.newFolder("StagingManagerTest");
+        mStagingManager = new StagingManager(mContext, null, mApexManager);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        if (mMockitoSession != null) {
+            mMockitoSession.finishMocking();
+        }
+    }
+
+    /**
+     * Tests that sessions committed later shouldn't cause earlier ones to fail the overlapping
+     * check.
+     */
+    @Test
+    public void checkNonOverlappingWithStagedSessions_laterSessionShouldNotFailEarlierOnes()
+            throws Exception {
+        // Create 2 sessions with overlapping packages
+        StagingManager.StagedSession session1 = createSession(111, "com.foo", 1);
+        StagingManager.StagedSession session2 = createSession(222, "com.foo", 2);
+
+        mStagingManager.createSession(session1);
+        mStagingManager.createSession(session2);
+        // Session1 should not fail in spite of the overlapping packages
+        mStagingManager.checkNonOverlappingWithStagedSessions(session1);
+        // Session2 should fail due to overlapping packages
+        assertThrows(PackageManagerException.class,
+                () -> mStagingManager.checkNonOverlappingWithStagedSessions(session2));
+    }
+
+    @Test
+    public void restoreSessions_nonParentSession_throwsIAE() throws Exception {
+        FakeStagedSession session = new FakeStagedSession(239);
+        session.setParentSessionId(1543);
+
+        assertThrows(IllegalArgumentException.class,
+                () -> mStagingManager.restoreSessions(Arrays.asList(session), false));
+    }
+
+    @Test
+    public void restoreSessions_nonCommittedSession_throwsIAE() throws Exception {
+        FakeStagedSession session = new FakeStagedSession(239);
+
+        assertThrows(IllegalArgumentException.class,
+                () -> mStagingManager.restoreSessions(Arrays.asList(session), false));
+    }
+
+    @Test
+    public void restoreSessions_terminalSession_throwsIAE() throws Exception {
+        FakeStagedSession session = new FakeStagedSession(239);
+        session.setCommitted(true);
+        session.setSessionApplied();
+
+        assertThrows(IllegalArgumentException.class,
+                () -> mStagingManager.restoreSessions(Arrays.asList(session), false));
+    }
+
+    @Test
+    public void restoreSessions_deviceUpgrading_failsAllSessions() throws Exception {
+        FakeStagedSession session1 = new FakeStagedSession(37);
+        session1.setCommitted(true);
+        FakeStagedSession session2 = new FakeStagedSession(57);
+        session2.setCommitted(true);
+
+        mStagingManager.restoreSessions(Arrays.asList(session1, session2), true);
+
+        assertThat(session1.getErrorCode()).isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(session1.getErrorMessage()).isEqualTo("Build fingerprint has changed");
+
+        assertThat(session2.getErrorCode()).isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(session2.getErrorMessage()).isEqualTo("Build fingerprint has changed");
+    }
+
+    @Test
+    public void restoreSessions_multipleSessions_deviceWithoutFsCheckpointSupport_throwISE()
+            throws Exception {
+        FakeStagedSession session1 = new FakeStagedSession(37);
+        session1.setCommitted(true);
+        FakeStagedSession session2 = new FakeStagedSession(57);
+        session2.setCommitted(true);
+
+        when(mStorageManager.supportsCheckpoint()).thenReturn(false);
+
+        assertThrows(IllegalStateException.class,
+                () -> mStagingManager.restoreSessions(Arrays.asList(session1, session2), false));
+    }
+
+    @Test
+    public void restoreSessions_handlesDestroyedAndNotReadySessions() throws Exception {
+        FakeStagedSession destroyedApkSession = new FakeStagedSession(23);
+        destroyedApkSession.setCommitted(true);
+        destroyedApkSession.setDestroyed(true);
+
+        FakeStagedSession destroyedApexSession = new FakeStagedSession(37);
+        destroyedApexSession.setCommitted(true);
+        destroyedApexSession.setDestroyed(true);
+        destroyedApexSession.setIsApex(true);
+
+        FakeStagedSession nonReadyApkSession = new FakeStagedSession(57);
+        nonReadyApkSession.setCommitted(true);
+
+        FakeStagedSession nonReadyApexSession = new FakeStagedSession(73);
+        nonReadyApexSession.setCommitted(true);
+        nonReadyApexSession.setIsApex(true);
+
+        FakeStagedSession destroyedNonReadySession = new FakeStagedSession(101);
+        destroyedNonReadySession.setCommitted(true);
+        destroyedNonReadySession.setDestroyed(true);
+
+        FakeStagedSession regularApkSession = new FakeStagedSession(239);
+        regularApkSession.setCommitted(true);
+        regularApkSession.setSessionReady();
+
+        List<StagingManager.StagedSession> sessions = new ArrayList<>();
+        sessions.add(destroyedApkSession);
+        sessions.add(destroyedApexSession);
+        sessions.add(nonReadyApkSession);
+        sessions.add(nonReadyApexSession);
+        sessions.add(destroyedNonReadySession);
+        sessions.add(regularApkSession);
+
+        mStagingManager.restoreSessions(sessions, false);
+
+        assertThat(sessions).containsExactly(regularApkSession);
+        assertThat(destroyedApkSession.isDestroyed()).isTrue();
+        assertThat(destroyedApexSession.isDestroyed()).isTrue();
+        assertThat(destroyedNonReadySession.isDestroyed()).isTrue();
+
+        mStagingManager.onBootCompletedBroadcastReceived();
+        assertThat(nonReadyApkSession.hasPreRebootVerificationStarted()).isTrue();
+        assertThat(nonReadyApexSession.hasPreRebootVerificationStarted()).isTrue();
+    }
+
+    @Test
+    public void restoreSessions_unknownApexSession_failsAllSessions() throws Exception {
+        FakeStagedSession apkSession = new FakeStagedSession(239);
+        apkSession.setCommitted(true);
+        apkSession.setSessionReady();
+
+        FakeStagedSession apexSession = new FakeStagedSession(1543);
+        apexSession.setCommitted(true);
+        apexSession.setIsApex(true);
+        apexSession.setSessionReady();
+
+        List<StagingManager.StagedSession> sessions = new ArrayList<>();
+        sessions.add(apkSession);
+        sessions.add(apexSession);
+
+        when(mApexManager.getSessions()).thenReturn(new SparseArray<>());
+        mStagingManager.restoreSessions(sessions, false);
+
+        // Validate checkpoint wasn't aborted.
+        verify(mStorageManager, never()).abortChanges(eq("abort-staged-install"), eq(false));
+
+        assertThat(apexSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apexSession.getErrorMessage()).isEqualTo("apexd did not know anything about a "
+                + "staged session supposed to be activated");
+
+        assertThat(apkSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apkSession.getErrorMessage()).isEqualTo("Another apex session failed");
+    }
+
+    @Test
+    public void restoreSessions_failedApexSessions_failsAllSessions() throws Exception {
+        FakeStagedSession apkSession = new FakeStagedSession(239);
+        apkSession.setCommitted(true);
+        apkSession.setSessionReady();
+
+        FakeStagedSession apexSession1 = new FakeStagedSession(1543);
+        apexSession1.setCommitted(true);
+        apexSession1.setIsApex(true);
+        apexSession1.setSessionReady();
+
+        FakeStagedSession apexSession2 = new FakeStagedSession(101);
+        apexSession2.setCommitted(true);
+        apexSession2.setIsApex(true);
+        apexSession2.setSessionReady();
+
+        FakeStagedSession apexSession3 = new FakeStagedSession(57);
+        apexSession3.setCommitted(true);
+        apexSession3.setIsApex(true);
+        apexSession3.setSessionReady();
+
+        ApexSessionInfo activationFailed = new ApexSessionInfo();
+        activationFailed.sessionId = 1543;
+        activationFailed.isActivationFailed = true;
+
+        ApexSessionInfo staged = new ApexSessionInfo();
+        staged.sessionId = 101;
+        staged.isStaged = true;
+
+        SparseArray<ApexSessionInfo> apexdSessions = new SparseArray<>();
+        apexdSessions.put(1543, activationFailed);
+        apexdSessions.put(101, staged);
+        when(mApexManager.getSessions()).thenReturn(apexdSessions);
+
+        List<StagingManager.StagedSession> sessions = new ArrayList<>();
+        sessions.add(apkSession);
+        sessions.add(apexSession1);
+        sessions.add(apexSession2);
+        sessions.add(apexSession3);
+
+        mStagingManager.restoreSessions(sessions, false);
+
+        // Validate checkpoint wasn't aborted.
+        verify(mStorageManager, never()).abortChanges(eq("abort-staged-install"), eq(false));
+
+        assertThat(apexSession1.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apexSession1.getErrorMessage()).isEqualTo("APEX activation failed. Check logcat "
+                + "messages from apexd for more information.");
+
+        assertThat(apexSession2.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apexSession2.getErrorMessage()).isEqualTo("Staged session 101 at boot didn't "
+                + "activate nor fail. Marking it as failed anyway.");
+
+        assertThat(apexSession3.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apexSession3.getErrorMessage()).isEqualTo("apexd did not know anything about a "
+                + "staged session supposed to be activated");
+
+        assertThat(apkSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apkSession.getErrorMessage()).isEqualTo("Another apex session failed");
+    }
+
+    @Test
+    public void restoreSessions_stagedApexSession_failsAllSessions() throws Exception {
+        FakeStagedSession apkSession = new FakeStagedSession(239);
+        apkSession.setCommitted(true);
+        apkSession.setSessionReady();
+
+        FakeStagedSession apexSession = new FakeStagedSession(1543);
+        apexSession.setCommitted(true);
+        apexSession.setIsApex(true);
+        apexSession.setSessionReady();
+
+        ApexSessionInfo staged = new ApexSessionInfo();
+        staged.sessionId = 1543;
+        staged.isStaged = true;
+
+        SparseArray<ApexSessionInfo> apexdSessions = new SparseArray<>();
+        apexdSessions.put(1543, staged);
+        when(mApexManager.getSessions()).thenReturn(apexdSessions);
+
+        List<StagingManager.StagedSession> sessions = new ArrayList<>();
+        sessions.add(apkSession);
+        sessions.add(apexSession);
+
+        mStagingManager.restoreSessions(sessions, false);
+
+        // Validate checkpoint wasn't aborted.
+        verify(mStorageManager, never()).abortChanges(eq("abort-staged-install"), eq(false));
+
+        assertThat(apexSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apexSession.getErrorMessage()).isEqualTo("Staged session 1543 at boot didn't "
+                + "activate nor fail. Marking it as failed anyway.");
+
+        assertThat(apkSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apkSession.getErrorMessage()).isEqualTo("Another apex session failed");
+    }
+
+    @Test
+    public void restoreSessions_failedAndActivatedApexSessions_abortsCheckpoint() throws Exception {
+        FakeStagedSession apkSession = new FakeStagedSession(239);
+        apkSession.setCommitted(true);
+        apkSession.setSessionReady();
+
+        FakeStagedSession apexSession1 = new FakeStagedSession(1543);
+        apexSession1.setCommitted(true);
+        apexSession1.setIsApex(true);
+        apexSession1.setSessionReady();
+
+        FakeStagedSession apexSession2 = new FakeStagedSession(101);
+        apexSession2.setCommitted(true);
+        apexSession2.setIsApex(true);
+        apexSession2.setSessionReady();
+
+        FakeStagedSession apexSession3 = new FakeStagedSession(57);
+        apexSession3.setCommitted(true);
+        apexSession3.setIsApex(true);
+        apexSession3.setSessionReady();
+
+        FakeStagedSession apexSession4 = new FakeStagedSession(37);
+        apexSession4.setCommitted(true);
+        apexSession4.setIsApex(true);
+        apexSession4.setSessionReady();
+
+        ApexSessionInfo activationFailed = new ApexSessionInfo();
+        activationFailed.sessionId = 1543;
+        activationFailed.isActivationFailed = true;
+
+        ApexSessionInfo activated = new ApexSessionInfo();
+        activated.sessionId = 101;
+        activated.isActivated = true;
+
+        ApexSessionInfo staged = new ApexSessionInfo();
+        staged.sessionId = 57;
+        staged.isActivationFailed = true;
+
+        SparseArray<ApexSessionInfo> apexdSessions = new SparseArray<>();
+        apexdSessions.put(1543, activationFailed);
+        apexdSessions.put(101, activated);
+        apexdSessions.put(57, staged);
+        when(mApexManager.getSessions()).thenReturn(apexdSessions);
+
+        List<StagingManager.StagedSession> sessions = new ArrayList<>();
+        sessions.add(apkSession);
+        sessions.add(apexSession1);
+        sessions.add(apexSession2);
+        sessions.add(apexSession3);
+        sessions.add(apexSession4);
+
+        mStagingManager.restoreSessions(sessions, false);
+
+        // Validate checkpoint was aborted.
+        verify(mStorageManager, times(1)).abortChanges(eq("abort-staged-install"), eq(false));
+    }
+
+    @Test
+    public void restoreSessions_apexSessionInImpossibleState_failsAllSessions() throws Exception {
+        FakeStagedSession apkSession = new FakeStagedSession(239);
+        apkSession.setCommitted(true);
+        apkSession.setSessionReady();
+
+        FakeStagedSession apexSession = new FakeStagedSession(1543);
+        apexSession.setCommitted(true);
+        apexSession.setIsApex(true);
+        apexSession.setSessionReady();
+
+        ApexSessionInfo impossible  = new ApexSessionInfo();
+        impossible.sessionId = 1543;
+
+        SparseArray<ApexSessionInfo> apexdSessions = new SparseArray<>();
+        apexdSessions.put(1543, impossible);
+        when(mApexManager.getSessions()).thenReturn(apexdSessions);
+
+        List<StagingManager.StagedSession> sessions = new ArrayList<>();
+        sessions.add(apkSession);
+        sessions.add(apexSession);
+
+        mStagingManager.restoreSessions(sessions, false);
+
+        // Validate checkpoint wasn't aborted.
+        verify(mStorageManager, never()).abortChanges(eq("abort-staged-install"), eq(false));
+
+        assertThat(apexSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apexSession.getErrorMessage()).isEqualTo("Impossible state");
+
+        assertThat(apkSession.getErrorCode())
+                .isEqualTo(SessionInfo.STAGED_SESSION_ACTIVATION_FAILED);
+        assertThat(apkSession.getErrorMessage()).isEqualTo("Another apex session failed");
+    }
+
+    private StagingManager.StagedSession createSession(int sessionId, String packageName,
+            long committedMillis) {
+        PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
+                PackageInstaller.SessionParams.MODE_FULL_INSTALL);
+        params.isStaged = true;
+
+        InstallSource installSource = InstallSource.create("testInstallInitiator",
+                "testInstallOriginator", "testInstaller", "testAttributionTag");
+
+        PackageInstallerSession session = new PackageInstallerSession(
+                /* callback */ null,
+                /* context */ null,
+                /* pm */ null,
+                /* sessionProvider */ null,
+                /* looper */ BackgroundThread.getHandler().getLooper(),
+                /* stagingManager */ null,
+                /* sessionId */ sessionId,
+                /* userId */ 456,
+                /* installerUid */ -1,
+                /* installSource */ installSource,
+                /* sessionParams */ params,
+                /* createdMillis */ 0L,
+                /* committedMillis */ committedMillis,
+                /* stageDir */ mTmpDir,
+                /* stageCid */ null,
+                /* files */ null,
+                /* checksums */ null,
+                /* prepared */ true,
+                /* committed */ true,
+                /* destroyed */ false,
+                /* sealed */ false,  // Setting to true would trigger some PM logic.
+                /* childSessionIds */ null,
+                /* parentSessionId */ -1,
+                /* isReady */ false,
+                /* isFailed */ false,
+                /* isApplied */false,
+                /* stagedSessionErrorCode */ PackageInstaller.SessionInfo.STAGED_SESSION_NO_ERROR,
+                /* stagedSessionErrorMessage */ "no error");
+
+        StagingManager.StagedSession stagedSession = spy(session.mStagedSession);
+        doReturn(packageName).when(stagedSession).getPackageName();
+        doAnswer(invocation -> {
+            Predicate<StagingManager.StagedSession> filter = invocation.getArgument(0);
+            return filter.test(stagedSession);
+        }).when(stagedSession).sessionContains(any());
+        return stagedSession;
+    }
+
+    private static final class FakeStagedSession implements StagingManager.StagedSession {
+        private final int mSessionId;
+        private boolean mIsApex = false;
+        private boolean mIsCommitted = false;
+        private boolean mIsReady = false;
+        private boolean mIsApplied = false;
+        private boolean mIsFailed = false;
+        private @StagedSessionErrorCode int mErrorCode = -1;
+        private String mErrorMessage;
+        private boolean mIsDestroyed = false;
+        private int mParentSessionId = -1;
+        private String mPackageName;
+        private boolean mIsAbandonded = false;
+        private boolean mPreRebootVerificationStarted = false;
+        private final List<StagingManager.StagedSession> mChildSessions = new ArrayList<>();
+
+        private FakeStagedSession(int sessionId) {
+            mSessionId = sessionId;
+        }
+
+        private void setParentSessionId(int parentSessionId) {
+            mParentSessionId = parentSessionId;
+        }
+
+        private void setCommitted(boolean isCommitted) {
+            mIsCommitted = isCommitted;
+        }
+
+        private void setIsApex(boolean isApex) {
+            mIsApex = isApex;
+        }
+
+        private void setDestroyed(boolean isDestroyed) {
+            mIsDestroyed = isDestroyed;
+        }
+
+        private void setPackageName(String packageName) {
+            mPackageName = packageName;
+        }
+
+        private boolean isAbandonded() {
+            return mIsAbandonded;
+        }
+
+        private boolean hasPreRebootVerificationStarted() {
+            return mPreRebootVerificationStarted;
+        }
+
+        private FakeStagedSession addChildSession(FakeStagedSession session) {
+            mChildSessions.add(session);
+            session.setParentSessionId(sessionId());
+            return this;
+        }
+
+        private @StagedSessionErrorCode int getErrorCode() {
+            return mErrorCode;
+        }
+
+        private String getErrorMessage() {
+            return mErrorMessage;
+        }
+
+        @Override
+        public boolean isMultiPackage() {
+            return !mChildSessions.isEmpty();
+        }
+
+        @Override
+        public boolean isApexSession() {
+            return mIsApex;
+        }
+
+        @Override
+        public boolean isCommitted() {
+            return mIsCommitted;
+        }
+
+        @Override
+        public boolean isInTerminalState() {
+            return isSessionApplied() || isSessionFailed();
+        }
+
+        @Override
+        public boolean isDestroyed() {
+            return mIsDestroyed;
+        }
+
+        @Override
+        public boolean isSessionReady() {
+            return mIsReady;
+        }
+
+        @Override
+        public boolean isSessionApplied() {
+            return mIsApplied;
+        }
+
+        @Override
+        public boolean isSessionFailed() {
+            return mIsFailed;
+        }
+
+        @Override
+        public List<StagingManager.StagedSession> getChildSessions() {
+            return mChildSessions;
+        }
+
+        @Override
+        public String getPackageName() {
+            return mPackageName;
+        }
+
+        @Override
+        public int getParentSessionId() {
+            return mParentSessionId;
+        }
+
+        @Override
+        public int sessionId() {
+            return mSessionId;
+        }
+
+        @Override
+        public PackageInstaller.SessionParams sessionParams() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public boolean sessionContains(Predicate<StagingManager.StagedSession> filter) {
+            return filter.test(this);
+        }
+
+        @Override
+        public boolean containsApkSession() {
+            Preconditions.checkState(!hasParentSessionId(), "Child session");
+            if (!isMultiPackage()) {
+                return !isApexSession();
+            }
+            for (StagingManager.StagedSession session : mChildSessions) {
+                if (!session.isApexSession()) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        @Override
+        public boolean containsApexSession() {
+            Preconditions.checkState(!hasParentSessionId(), "Child session");
+            if (!isMultiPackage()) {
+                return isApexSession();
+            }
+            for (StagingManager.StagedSession session : mChildSessions) {
+                if (session.isApexSession()) {
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        @Override
+        public void setSessionReady() {
+            mIsReady = true;
+        }
+
+        @Override
+        public void setSessionFailed(@StagedSessionErrorCode int errorCode, String errorMessage) {
+            Preconditions.checkState(!mIsApplied, "Already marked as applied");
+            mIsFailed = true;
+            mErrorCode = errorCode;
+            mErrorMessage = errorMessage;
+        }
+
+        @Override
+        public void setSessionApplied() {
+            Preconditions.checkState(!mIsFailed, "Already marked as failed");
+            mIsApplied = true;
+        }
+
+        @Override
+        public void installSession(IntentSender statusReceiver) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public boolean hasParentSessionId() {
+            return mParentSessionId != -1;
+        }
+
+        @Override
+        public long getCommittedMillis() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void abandon() {
+            mIsAbandonded = true;
+        }
+
+        @Override
+        public boolean notifyStartPreRebootVerification() {
+            mPreRebootVerificationStarted = true;
+            // TODO(ioffe): change to true when tests for pre-reboot verification are added.
+            return false;
+        }
+
+        @Override
+        public void notifyEndPreRebootVerification() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void verifySession() {
+            throw new UnsupportedOperationException();
+        }
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/pm/StagingManagerTest.java b/services/tests/servicestests/src/com/android/server/pm/StagingManagerTest.java
deleted file mode 100644
index 79935c23774f..000000000000
--- a/services/tests/servicestests/src/com/android/server/pm/StagingManagerTest.java
+++ /dev/null
@@ -1,135 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.server.pm;
-
-import android.content.Context;
-import android.content.pm.PackageInstaller;
-import android.os.storage.StorageManager;
-import android.platform.test.annotations.Presubmit;
-
-import com.android.internal.os.BackgroundThread;
-
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.TemporaryFolder;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import org.mockito.Mockito;
-import org.mockito.MockitoAnnotations;
-
-import java.io.File;
-import java.util.function.Predicate;
-
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.when;
-import static org.testng.Assert.assertThrows;
-
-@Presubmit
-@RunWith(JUnit4.class)
-public class StagingManagerTest {
-    @Rule
-    public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
-
-    private File mTmpDir;
-    private StagingManager mStagingManager;
-
-    @Before
-    public void setup() throws Exception {
-        MockitoAnnotations.initMocks(this);
-        StorageManager storageManager = Mockito.mock(StorageManager.class);
-        Context context = Mockito.mock(Context.class);
-        when(storageManager.isCheckpointSupported()).thenReturn(true);
-        when(context.getSystemService(eq(Context.POWER_SERVICE))).thenReturn(null);
-        when(context.getSystemService(eq(Context.STORAGE_SERVICE))).thenReturn(storageManager);
-
-        mTmpDir = mTemporaryFolder.newFolder("StagingManagerTest");
-        mStagingManager = new StagingManager(context, null);
-    }
-
-    /**
-     * Tests that sessions committed later shouldn't cause earlier ones to fail the overlapping
-     * check.
-     */
-    @Test
-    public void checkNonOverlappingWithStagedSessions_laterSessionShouldNotFailEarlierOnes()
-            throws Exception {
-        // Create 2 sessions with overlapping packages
-        StagingManager.StagedSession session1 = createSession(111, "com.foo", 1);
-        StagingManager.StagedSession session2 = createSession(222, "com.foo", 2);
-
-        mStagingManager.createSession(session1);
-        mStagingManager.createSession(session2);
-        // Session1 should not fail in spite of the overlapping packages
-        mStagingManager.checkNonOverlappingWithStagedSessions(session1);
-        // Session2 should fail due to overlapping packages
-        assertThrows(PackageManagerException.class,
-                () -> mStagingManager.checkNonOverlappingWithStagedSessions(session2));
-    }
-
-    private StagingManager.StagedSession createSession(int sessionId, String packageName,
-            long committedMillis) {
-        PackageInstaller.SessionParams params = new PackageInstaller.SessionParams(
-                PackageInstaller.SessionParams.MODE_FULL_INSTALL);
-        params.isStaged = true;
-
-        InstallSource installSource = InstallSource.create("testInstallInitiator",
-                "testInstallOriginator", "testInstaller", "testAttributionTag");
-
-        PackageInstallerSession session = new PackageInstallerSession(
-                /* callback */ null,
-                /* context */ null,
-                /* pm */ null,
-                /* sessionProvider */ null,
-                /* looper */ BackgroundThread.getHandler().getLooper(),
-                /* stagingManager */ null,
-                /* sessionId */ sessionId,
-                /* userId */ 456,
-                /* installerUid */ -1,
-                /* installSource */ installSource,
-                /* sessionParams */ params,
-                /* createdMillis */ 0L,
-                /* committedMillis */ committedMillis,
-                /* stageDir */ mTmpDir,
-                /* stageCid */ null,
-                /* files */ null,
-                /* checksums */ null,
-                /* prepared */ true,
-                /* committed */ true,
-                /* destroyed */ false,
-                /* sealed */ false,  // Setting to true would trigger some PM logic.
-                /* childSessionIds */ null,
-                /* parentSessionId */ -1,
-                /* isReady */ false,
-                /* isFailed */ false,
-                /* isApplied */false,
-                /* stagedSessionErrorCode */ PackageInstaller.SessionInfo.STAGED_SESSION_NO_ERROR,
-                /* stagedSessionErrorMessage */ "no error");
-
-        StagingManager.StagedSession stagedSession = spy(session.mStagedSession);
-        doReturn(packageName).when(stagedSession).getPackageName();
-        doAnswer(invocation -> {
-            Predicate<StagingManager.StagedSession> filter = invocation.getArgument(0);
-            return filter.test(stagedSession);
-        }).when(stagedSession).sessionContains(any());
-        return stagedSession;
-    }
-}