Allow handling of all pending staged rollback sessions before rebooting

When there is an unattributable native service crash during boot, we
rollback all packages available for rollback. In the current
implementation, we reboot immediately when a staged rollback session
becomes ready, even though there can be other staged rollback sessions
which are still being handled.

In this CL, the logic is changed to allow the health observer wait for
all staged sessions to be handled before rebooting.

Bug: 141843321
Test: atest StagedRollbackTest#testNativeWatchdogTriggersRollbackForAll
Test: atest StagedRollbackTest
Change-Id: Ib2aaee22868e5502b68847aa7f64e428070f7e3c

services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java

@@ -280,7 +280,6 @@ public final class RollbackPackageHealthObserver implements PackageHealthObserve
                             .WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_BOOT_TRIGGERED,
                             WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_REASON__REASON_UNKNOWN,
                             "");
-                    mContext.getSystemService(PowerManager.class).reboot("Rollback staged install");
                 } else if (sessionInfo.isStagedSessionFailed()
                         && markStagedSessionHandled(rollbackId)) {
                     logEvent(moduleMetadataPackage,
@@ -291,6 +290,11 @@ public final class RollbackPackageHealthObserver implements PackageHealthObserve
                 }
             }
         }
+
+        // Wait for all pending staged sessions to get handled before rebooting.
+        if (isPendingStagedSessionsEmpty()) {
+            mContext.getSystemService(PowerManager.class).reboot("Rollback staged install");
+        }
     }
 
     /**
@@ -303,6 +307,16 @@ public final class RollbackPackageHealthObserver implements PackageHealthObserve
         }
     }
 
+    /**
+     * Returns {@code true} if all pending staged rollback sessions were marked as handled,
+     * {@code false} if there is any left.
+     */
+    private boolean isPendingStagedSessionsEmpty() {
+        synchronized (mPendingStagedRollbackIds) {
+            return mPendingStagedRollbackIds.isEmpty();
+        }
+    }
+
     private void saveLastStagedRollbackId(int stagedRollbackId) {
         try {
             FileOutputStream fos = new FileOutputStream(mLastStagedRollbackIdFile);
@@ -414,6 +428,9 @@ public final class RollbackPackageHealthObserver implements PackageHealthObserve
                             reasonToLog, failedPackageToLog);
                 }
             } else {
+                if (rollback.isStaged()) {
+                    markStagedSessionHandled(rollback.getRollbackId());
+                }
                 logEvent(logPackage,
                         StatsLog.WATCHDOG_ROLLBACK_OCCURRED__ROLLBACK_TYPE__ROLLBACK_FAILURE,
                         reasonToLog, failedPackageToLog);
@@ -431,6 +448,16 @@ public final class RollbackPackageHealthObserver implements PackageHealthObserve
         RollbackManager rollbackManager = mContext.getSystemService(RollbackManager.class);
         List<RollbackInfo> rollbacks = rollbackManager.getAvailableRollbacks();
 
+        // Add all rollback ids to mPendingStagedRollbackIds, so that we do not reboot before all
+        // pending staged rollbacks are handled.
+        synchronized (mPendingStagedRollbackIds) {
+            for (RollbackInfo rollback : rollbacks) {
+                if (rollback.isStaged()) {
+                    mPendingStagedRollbackIds.add(rollback.getRollbackId());
+                }
+            }
+        }
+
         for (RollbackInfo rollback : rollbacks) {
             VersionedPackage sample = rollback.getPackages().get(0).getVersionRolledBackFrom();
             rollbackPackage(rollback, sample, PackageWatchdog.FAILURE_REASON_NATIVE_CRASH);

tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java

@@ -23,12 +23,14 @@ import static com.google.common.truth.Truth.assertThat;
 
 import android.Manifest;
 import android.content.ComponentName;
+import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageManager;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
 import android.os.ParcelFileDescriptor;
+import android.os.storage.StorageManager;
 import android.provider.DeviceConfig;
 
 import androidx.test.platform.app.InstrumentationRegistry;
@@ -498,4 +500,11 @@ public class StagedRollbackTest {
                 .executeShellCommand(cmd);
         IoUtils.closeQuietly(pfd);
     }
+
+    @Test
+    public void isCheckpointSupported() {
+        Context context = InstrumentationRegistry.getInstrumentation().getContext();
+        StorageManager sm = (StorageManager) context.getSystemService(Context.STORAGE_SERVICE);
+        assertThat(sm.isCheckpointSupported()).isTrue();
+    }
 }

tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java

@@ -17,6 +17,7 @@
 package com.android.tests.rollback.host;
 
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeTrue;
 import static org.testng.Assert.assertThrows;
 
 import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
@@ -103,6 +104,9 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
 
     @Test
     public void testNativeWatchdogTriggersRollbackForAll() throws Exception {
+        // This test requires committing multiple staged rollbacks
+        assumeTrue(isCheckpointSupported());
+
         // Install a package with rollback enabled.
         runPhase("testNativeWatchdogTriggersRollbackForAll_Phase1");
         getDevice().reboot();
@@ -233,4 +237,13 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
         // Find the NetworkStack path (can be NetworkStack.apk or NetworkStackNext.apk)
         return getDevice().executeShellCommand("ls /system/priv-app/NetworkStack*/*.apk");
     }
+
+    private boolean isCheckpointSupported() throws Exception {
+        try {
+            runPhase("isCheckpointSupported");
+            return true;
+        } catch (AssertionError ignore) {
+            return false;
+        }
+    }
 }