Merge changes I5dcaabaa,I998cf688

* changes:
  Make rollback-app support --staged-ready-timeout flag
  Simplify flags used to wait for staged session ready

3 files changed, 104 insertions, 68 deletions

diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index 39b320318302..b571a9c5ce13 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -104,6 +104,7 @@ import android.util.SparseArray;
 import com.android.internal.content.PackageHelper;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.IndentingPrintWriter;
+import com.android.internal.util.Preconditions;
 import com.android.server.LocalServices;
 import com.android.server.SystemConfig;
 import com.android.server.pm.PackageManagerShellCommandDataLoader.Metadata;
@@ -139,7 +140,7 @@ class PackageManagerShellCommand extends ShellCommand {
     private static final String STDIN_PATH = "-";
     /** Path where ART profiles snapshots are dumped for the shell user */
     private final static String ART_PROFILE_SNAPSHOT_DEBUG_LOCATION = "/data/misc/profman/";
-    private static final int DEFAULT_WAIT_MS = 60 * 1000;
+    private static final int DEFAULT_STAGED_READY_TIMEOUT_MS = 60 * 1000;
     private static final String TAG = "PackageManagerShellCommand";
 
     final IPackageManager mInterface;
@@ -455,9 +456,20 @@ class PackageManagerShellCommand extends ShellCommand {
         return 1;
     }
 
-    private int runRollbackApp() {
+    private int runRollbackApp() throws RemoteException {
         final PrintWriter pw = getOutPrintWriter();
 
+        String opt;
+        long stagedReadyTimeoutMs = DEFAULT_STAGED_READY_TIMEOUT_MS;
+        while ((opt = getNextOption()) != null) {
+            switch (opt) {
+                case "--staged-ready-timeout":
+                    stagedReadyTimeoutMs = Long.parseLong(getNextArgRequired());
+                    break;
+                default:
+                    throw new IllegalArgumentException("Unknown option: " + opt);
+            }
+        }
         final String packageName = getNextArgRequired();
         if (packageName == null) {
             pw.println("Error: package name not specified");
@@ -465,11 +477,10 @@ class PackageManagerShellCommand extends ShellCommand {
         }
 
         final LocalIntentReceiver receiver = new LocalIntentReceiver();
+        RollbackInfo rollback = null;
         try {
             IRollbackManager rm = IRollbackManager.Stub.asInterface(
                     ServiceManager.getService(Context.ROLLBACK_SERVICE));
-
-            RollbackInfo rollback = null;
             for (RollbackInfo r : (List<RollbackInfo>) rm.getAvailableRollbacks().getList()) {
                 for (PackageRollbackInfo info : r.getPackages()) {
                     if (packageName.equals(info.getPackageName())) {
@@ -494,14 +505,21 @@ class PackageManagerShellCommand extends ShellCommand {
         final Intent result = receiver.getResult();
         final int status = result.getIntExtra(RollbackManager.EXTRA_STATUS,
                 RollbackManager.STATUS_FAILURE);
-        if (status == RollbackManager.STATUS_SUCCESS) {
-            pw.println("Success");
-            return 0;
-        } else {
+
+        if (status != RollbackManager.STATUS_SUCCESS) {
             pw.println("Failure ["
                     + result.getStringExtra(RollbackManager.EXTRA_STATUS_MESSAGE) + "]");
             return 1;
         }
+
+        if (rollback.isStaged() && stagedReadyTimeoutMs > 0) {
+            final int committedSessionId = rollback.getCommittedSessionId();
+            return doWaitForStagedSessionReady(committedSessionId, stagedReadyTimeoutMs, pw);
+        }
+
+        pw.println("Success");
+        return 0;
+
     }
 
     private void setParamsSize(InstallParams params, List<String> inPaths) {
@@ -1304,11 +1322,12 @@ class PackageManagerShellCommand extends ShellCommand {
             }
             abandonSession = false;
 
-            if (!params.sessionParams.isStaged || !params.mWaitForStagedSessionReady) {
-                pw.println("Success");
-                return 0;
+            if (params.sessionParams.isStaged && params.stagedReadyTimeoutMs > 0) {
+                return doWaitForStagedSessionReady(sessionId, params.stagedReadyTimeoutMs, pw);
             }
-            return doWaitForStagedSessionRead(sessionId, params.timeoutMs, pw);
+
+            pw.println("Success");
+            return 0;
         } finally {
             if (abandonSession) {
                 try {
@@ -1319,11 +1338,9 @@ class PackageManagerShellCommand extends ShellCommand {
         }
     }
 
-    private int doWaitForStagedSessionRead(int sessionId, long timeoutMs, PrintWriter pw)
+    private int doWaitForStagedSessionReady(int sessionId, long timeoutMs, PrintWriter pw)
               throws RemoteException {
-        if (timeoutMs <= 0) {
-            timeoutMs = DEFAULT_WAIT_MS;
-        }
+        Preconditions.checkArgument(timeoutMs > 0);
         PackageInstaller.SessionInfo si = mInterface.getPackageInstaller()
                 .getSessionInfo(sessionId);
         if (si == null) {
@@ -1373,25 +1390,14 @@ class PackageManagerShellCommand extends ShellCommand {
     private int runInstallCommit() throws RemoteException {
         final PrintWriter pw = getOutPrintWriter();
         String opt;
-        boolean waitForStagedSessionReady = true;
-        long timeoutMs = -1;
+        long stagedReadyTimeoutMs = DEFAULT_STAGED_READY_TIMEOUT_MS;
         while ((opt = getNextOption()) != null) {
             switch (opt) {
-                case "--wait-for-staged-ready":
-                    waitForStagedSessionReady = true;
-                    // If there is only one remaining argument, then it represents the sessionId, we
-                    // shouldn't try to parse it as timeoutMs.
-                    if (getRemainingArgsCount() > 1) {
-                        try {
-                            timeoutMs = Long.parseLong(peekNextArg());
-                            getNextArg();
-                        } catch (NumberFormatException ignore) {
-                        }
-                    }
-                    break;
-                case "--no-wait":
-                    waitForStagedSessionReady = false;
+                case "--staged-ready-timeout":
+                    stagedReadyTimeoutMs = Long.parseLong(getNextArgRequired());
                     break;
+                default:
+                    throw new IllegalArgumentException("Unknown option: " + opt);
             }
         }
         final int sessionId = Integer.parseInt(getNextArg());
@@ -1400,11 +1406,11 @@ class PackageManagerShellCommand extends ShellCommand {
         }
         final PackageInstaller.SessionInfo si = mInterface.getPackageInstaller()
                 .getSessionInfo(sessionId);
-        if (si == null || !si.isStaged() || !waitForStagedSessionReady) {
-            pw.println("Success");
-            return 0;
+        if (si != null && si.isStaged() && stagedReadyTimeoutMs > 0) {
+            return doWaitForStagedSessionReady(sessionId, stagedReadyTimeoutMs, pw);
         }
-        return doWaitForStagedSessionRead(sessionId, timeoutMs, pw);
+        pw.println("Success");
+        return 0;
     }
 
     private int runInstallCreate() throws RemoteException {
@@ -2735,8 +2741,7 @@ class PackageManagerShellCommand extends ShellCommand {
         SessionParams sessionParams;
         String installerPackageName;
         int userId = UserHandle.USER_ALL;
-        boolean mWaitForStagedSessionReady = true;
-        long timeoutMs = DEFAULT_WAIT_MS;
+        long stagedReadyTimeoutMs = DEFAULT_STAGED_READY_TIMEOUT_MS;
     }
 
     private InstallParams makeInstallParams() {
@@ -2865,16 +2870,8 @@ class PackageManagerShellCommand extends ShellCommand {
                     }
                     sessionParams.installFlags |= PackageManager.INSTALL_ENABLE_ROLLBACK;
                     break;
-                case "--wait-for-staged-ready":
-                    params.mWaitForStagedSessionReady = true;
-                    try {
-                        params.timeoutMs = Long.parseLong(peekNextArg());
-                        getNextArg();
-                    } catch (NumberFormatException ignore) {
-                    }
-                    break;
-                case "--no-wait":
-                    params.mWaitForStagedSessionReady = false;
+                case "--staged-ready-timeout":
+                    params.stagedReadyTimeoutMs = Long.parseLong(getNextArgRequired());
                     break;
                 case "--skip-verification":
                     sessionParams.installFlags |= PackageManager.INSTALL_DISABLE_VERIFICATION;
@@ -3597,7 +3594,7 @@ class PackageManagerShellCommand extends ShellCommand {
         pw.println("       [--preload] [--instant] [--full] [--dont-kill]");
         pw.println("       [--enable-rollback]");
         pw.println("       [--force-uuid internal|UUID] [--pkg PACKAGE] [-S BYTES]");
-        pw.println("       [--apex] [--wait-for-staged-ready TIMEOUT]");
+        pw.println("       [--apex] [--staged-ready-timeout TIMEOUT]");
         pw.println("       [PATH [SPLIT...]|-]");
         pw.println("    Install an application.  Must provide the apk data to install, either as");
         pw.println("    file path(s) or '-' to read from stdin.  Options are:");
@@ -3625,9 +3622,11 @@ class PackageManagerShellCommand extends ShellCommand {
         pw.println("          3=device setup, 4=user request");
         pw.println("      --force-uuid: force install on to disk volume with given UUID");
         pw.println("      --apex: install an .apex file, not an .apk");
-        pw.println("      --wait-for-staged-ready: when performing staged install, wait TIMEOUT");
-        pw.println("          ms for pre-reboot verification to complete. If TIMEOUT is not");
-        pw.println("          specified it will wait for " + DEFAULT_WAIT_MS + " milliseconds.");
+        pw.println("      --staged-ready-timeout: By default, staged sessions wait "
+                + DEFAULT_STAGED_READY_TIMEOUT_MS);
+        pw.println("          milliseconds for pre-reboot verification to complete when");
+        pw.println("          performing staged install. This flag is used to alter the waiting");
+        pw.println("          time. You can skip the waiting time by specifying a TIMEOUT of '0'");
         pw.println("");
         pw.println("  install-existing [--user USER_ID|all|current]");
         pw.println("       [--instant] [--full] [--wait] [--restrict-permissions] PACKAGE");
diff --git a/tests/StagedInstallTest/Android.bp b/tests/StagedInstallTest/Android.bp
index 0e7a049f5faa..b001fe105bc1 100644
--- a/tests/StagedInstallTest/Android.bp
+++ b/tests/StagedInstallTest/Android.bp
@@ -23,7 +23,7 @@ android_test_helper_app {
 java_test_host {
     name: "StagedInstallInternalTest",
     srcs: ["src/**/*.java"],
-    libs: ["tradefed"],
+    libs: ["tradefed", "cts-shim-host-lib"],
     static_libs: [
         "testng",
         "compatibility-tradefed",
diff --git a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
index ccd18dd25167..407c65ba4944 100644
--- a/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
+++ b/tests/StagedInstallTest/src/com/android/tests/stagedinstallinternal/host/StagedInstallInternalTest.java
@@ -16,6 +16,8 @@
 
 package com.android.tests.stagedinstallinternal.host;
 
+import static com.android.cts.shim.lib.ShimPackage.SHIM_APEX_PACKAGE_NAME;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertTrue;
@@ -93,43 +95,78 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         runPhase("testSystemServerRestartDoesNotAffectStagedSessions_Verify");
     }
 
+    // Test waiting time for staged session to be ready using adb staged install can be altered
     @Test
-    public void testAdbStagedInstallWaitForReadyFlagWorks() throws Exception {
+    public void testAdbStagdReadyTimeoutFlagWorks() throws Exception {
         assumeTrue("Device does not support updating APEX",
                 mHostUtils.isApexUpdateSupported());
 
-        File apexFile = mTestUtils.getTestFile(SHIM_V2);
-        String output = getDevice().executeAdbCommand("install", "--staged",
-                "--wait-for-staged-ready", "60000", apexFile.getAbsolutePath());
+        final File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        final String output = getDevice().executeAdbCommand("install", "--staged",
+                "--staged-ready-timeout", "60000", apexFile.getAbsolutePath());
         assertThat(output).contains("Reboot device to apply staged session");
-        String sessionId = getDevice().executeShellCommand(
+        final String sessionId = getDevice().executeShellCommand(
                 "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
         assertThat(sessionId).isNotEmpty();
     }
 
+    // Test adb staged installation wait for session to be ready by default
     @Test
-    public void testAdbStagedInstallNoWaitFlagWorks() throws Exception {
+    public void testAdbStagedInstallWaitsTillReadyByDefault() throws Exception {
         assumeTrue("Device does not support updating APEX",
                 mHostUtils.isApexUpdateSupported());
 
-        File apexFile = mTestUtils.getTestFile(SHIM_V2);
-        String output = getDevice().executeAdbCommand("install", "--staged",
-                "--no-wait", apexFile.getAbsolutePath());
+        final File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        final String output = getDevice().executeAdbCommand("install", "--staged",
+                apexFile.getAbsolutePath());
+        assertThat(output).contains("Reboot device to apply staged session");
+        final String sessionId = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
+        assertThat(sessionId).isNotEmpty();
+    }
+
+    // Test we can skip waiting for staged session to be ready
+    @Test
+    public void testAdbStagedReadyWaitCanBeSkipped() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mHostUtils.isApexUpdateSupported());
+
+        final File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        final String output = getDevice().executeAdbCommand("install", "--staged",
+                "--staged-ready-timeout", "0", apexFile.getAbsolutePath());
         assertThat(output).doesNotContain("Reboot device to apply staged session");
         assertThat(output).contains("Success");
-        String sessionId = getDevice().executeShellCommand(
+        final String sessionId = getDevice().executeShellCommand(
                 "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
         assertThat(sessionId).isEmpty();
     }
 
+    // Test rollback-app command waits for staged sessions to be ready
+    @Test
+    public void testAdbRollbackAppWaitsForStagedReady() throws Exception {
+        assumeTrue("Device does not support updating APEX",
+                mHostUtils.isApexUpdateSupported());
+
+        final File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        String output = getDevice().executeAdbCommand("install", "--staged",
+                "--enable-rollback", apexFile.getAbsolutePath());
+        assertThat(output).contains("Reboot device to apply staged session");
+        getDevice().reboot();
+        output = getDevice().executeShellCommand("pm rollback-app " + SHIM_APEX_PACKAGE_NAME);
+        assertThat(output).contains("Reboot device to apply staged session");
+        final String sessionId = getDevice().executeShellCommand(
+                "pm get-stagedsessions --only-ready --only-parent --only-sessionid").trim();
+        assertThat(sessionId).isNotEmpty();
+    }
+
     @Test
     public void testAdbInstallMultiPackageCommandWorks() throws Exception {
         assumeTrue("Device does not support updating APEX",
                 mHostUtils.isApexUpdateSupported());
 
-        File apexFile = mTestUtils.getTestFile(SHIM_V2);
-        File apkFile = mTestUtils.getTestFile(APK_A);
-        String output = getDevice().executeAdbCommand("install-multi-package",
+        final File apexFile = mTestUtils.getTestFile(SHIM_V2);
+        final File apkFile = mTestUtils.getTestFile(APK_A);
+        final String output = getDevice().executeAdbCommand("install-multi-package",
                 apexFile.getAbsolutePath(), apkFile.getAbsolutePath());
         assertThat(output).contains("Created parent session");
         assertThat(output).contains("Created child session");
@@ -154,10 +191,10 @@ public class StagedInstallInternalTest extends BaseHostJUnit4Test {
         getDevice().disableAdbRoot();
 
         // Wait for new system server process to start
-        long start = System.currentTimeMillis();
+        final long start = System.currentTimeMillis();
         long newStartTime = oldStartTime;
         while (System.currentTimeMillis() < start + SYSTEM_SERVER_TIMEOUT_MS) {
-            ProcessInfo newPs = getDevice().getProcessByName("system_server");
+            final ProcessInfo newPs = getDevice().getProcessByName("system_server");
             if (newPs != null) {
                 newStartTime = newPs.getStartTime();
                 if (newStartTime != oldStartTime) {