Preserve window for Activity#recreate if possible

To prevent a black screen if the activity calls recreate in
the foreground.

Bug: 133216672
Test: ActivityThreadTest#testRecreateActivity
      ActivityThreadTest#testHandleActivity_assetsChanged
Change-Id: I47d22895287f94cc77130767b67d4b0bdf82c3a3

3 files changed, 48 insertions, 25 deletions

diff --git a/core/java/android/app/ActivityThread.java b/core/java/android/app/ActivityThread.java
index b68194792db1..a9ff0e90a3ec 100644
--- a/core/java/android/app/ActivityThread.java
+++ b/core/java/android/app/ActivityThread.java
@@ -4971,15 +4971,8 @@ public final class ActivityThread extends ClientTransactionHandler {
     }
 
     private void relaunchAllActivities(boolean preserveWindows) {
-        for (Map.Entry<IBinder, ActivityClientRecord> entry : mActivities.entrySet()) {
-            final ActivityClientRecord r = entry.getValue();
-            // Schedule relaunch the activity if it is not a local object or finishing.
-            if (!r.activity.mFinished && !(r.token instanceof Binder)) {
-                if (preserveWindows && r.window != null) {
-                    r.mPreserveWindow = true;
-                }
-                scheduleRelaunchActivity(entry.getKey());
-            }
+        for (int i = mActivities.size() - 1; i >= 0; i--) {
+            scheduleRelaunchActivityIfPossible(mActivities.valueAt(i), preserveWindows);
         }
     }
 
@@ -5348,15 +5341,31 @@ public final class ActivityThread extends ClientTransactionHandler {
         }
     }
 
+    void scheduleRelaunchActivity(IBinder token) {
+        final ActivityClientRecord r = mActivities.get(token);
+        if (r != null) {
+            scheduleRelaunchActivityIfPossible(r, !r.stopped /* preserveWindow */);
+        }
+    }
+
     /**
      * Post a message to relaunch the activity. We do this instead of launching it immediately,
      * because this will destroy the activity from which it was called and interfere with the
      * lifecycle changes it was going through before. We need to make sure that we have finished
      * handling current transaction item before relaunching the activity.
      */
-    void scheduleRelaunchActivity(IBinder token) {
-        mH.removeMessages(H.RELAUNCH_ACTIVITY, token);
-        sendMessage(H.RELAUNCH_ACTIVITY, token);
+    private void scheduleRelaunchActivityIfPossible(@NonNull ActivityClientRecord r,
+            boolean preserveWindow) {
+        if (r.activity.mFinished || r.token instanceof Binder) {
+            // Do not schedule relaunch if the activity is finishing or not a local object (e.g.
+            // created by ActivtiyGroup that server side doesn't recognize it).
+            return;
+        }
+        if (preserveWindow && r.window != null) {
+            r.mPreserveWindow = true;
+        }
+        mH.removeMessages(H.RELAUNCH_ACTIVITY, r.token);
+        sendMessage(H.RELAUNCH_ACTIVITY, r.token);
     }
 
     /** Performs the activity relaunch locally vs. requesting from system-server. */
diff --git a/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java b/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java
index 400b05c09fa5..4fe68cd5a27a 100644
--- a/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java
+++ b/core/tests/coretests/src/android/app/activity/ActivityThreadTest.java
@@ -73,6 +73,7 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 /**
  * Test for verifying {@link android.app.ActivityThread} class.
@@ -173,29 +174,41 @@ public class ActivityThreadTest {
 
     @Test
     public void testHandleActivity_assetsChanged() {
+        relaunchActivityAndAssertPreserveWindow(activity -> {
+            // Relaunches all activities.
+            activity.getActivityThread().handleApplicationInfoChanged(
+                    activity.getApplicationInfo());
+        });
+    }
+
+    @Test
+    public void testRecreateActivity() {
+        relaunchActivityAndAssertPreserveWindow(Activity::recreate);
+    }
+
+    private void relaunchActivityAndAssertPreserveWindow(Consumer<Activity> relaunch) {
         final TestActivity activity = mActivityTestRule.launchActivity(new Intent());
+        final ActivityThread activityThread = activity.getActivityThread();
 
         final IBinder[] token = new IBinder[1];
         final View[] decorView = new View[1];
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
-            final ActivityThread activityThread = activity.getActivityThread();
-
             token[0] = activity.getActivityToken();
             decorView[0] = activity.getWindow().getDecorView();
 
-            // Relaunches all activities
-            activityThread.handleApplicationInfoChanged(activity.getApplicationInfo());
+            relaunch.accept(activity);
         });
 
         final View[] newDecorView = new View[1];
-        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
-            final ActivityThread activityThread = activity.getActivityThread();
+        final Activity[] newActivity = new Activity[1];
 
-            final Activity newActivity = activityThread.getActivity(token[0]);
-            newDecorView[0] = activity.getWindow().getDecorView();
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            newActivity[0] = activityThread.getActivity(token[0]);
+            newDecorView[0] = newActivity[0].getWindow().getDecorView();
         });
 
+        assertNotEquals("Activity must be relaunched", activity, newActivity[0]);
         assertEquals("Window must be preserved", decorView[0], newDecorView[0]);
     }
 
diff --git a/core/tests/mockingcoretests/src/android/app/activity/ActivityThreadClientTest.java b/core/tests/mockingcoretests/src/android/app/activity/ActivityThreadClientTest.java
index d266cdbccf79..46695d2764dd 100644
--- a/core/tests/mockingcoretests/src/android/app/activity/ActivityThreadClientTest.java
+++ b/core/tests/mockingcoretests/src/android/app/activity/ActivityThreadClientTest.java
@@ -26,15 +26,16 @@ import static android.app.servertransaction.ActivityLifecycleItem.ON_STOP;
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.after;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -52,7 +53,7 @@ import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.res.CompatibilityInfo;
 import android.content.res.Configuration;
-import android.os.Binder;
+import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
@@ -200,9 +201,9 @@ public class ActivityThreadClientTest {
     private void recreateAndVerifyNoRelaunch(ActivityThread activityThread, TestActivity activity) {
         clearInvocations(activityThread);
         getInstrumentation().runOnMainSync(() -> activity.recreate());
+        getInstrumentation().waitForIdleSync();
 
-        verify(activityThread, after(WAIT_TIMEOUT_MS).never())
-                .handleRelaunchActivity(any(), any());
+        verify(activityThread, never()).handleRelaunchActivity(any(), any());
     }
 
     private void recreateAndVerifyRelaunched(ActivityThread activityThread, TestActivity activity,
@@ -292,7 +293,7 @@ public class ActivityThreadClientTest {
             spyOn(packageInfo);
             doNothing().when(packageInfo).updateApplicationInfo(any(), any());
 
-            return new ActivityClientRecord(new Binder(), Intent.makeMainActivity(component),
+            return new ActivityClientRecord(mock(IBinder.class), Intent.makeMainActivity(component),
                     0 /* ident */, info, new Configuration(),
                     CompatibilityInfo.DEFAULT_COMPATIBILITY_INFO, null /* referrer */,
                     null /* voiceInteractor */, null /* state */, null /* persistentState */,