Apply transaction when starting back nav

The transaction was not applied anymore after the refactor and the
end of animation callback was not correctly set.

Bug: 228202811
Test: com.android.server.wm.BackNavigationControllerTests#backNavInfo_HomeWhenBackToLauncher

Change-Id: Id4635f2988a23ad7e214fb37739ba8dc2c394ee4
diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index eb16394..428a99c 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -703,12 +703,6 @@
       "group": "WM_DEBUG_STATES",
       "at": "com\/android\/server\/wm\/ActivityRecord.java"
     },
-    "-1427392850": {
-      "message": "WindowState: Setting back callback %s (priority: %d) (Client IWindow: %s). (WindowState: %s)",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/WindowState.java"
-    },
     "-1427184084": {
       "message": "addWindow: New client %s: window=%s Callers=%s",
       "level": "VERBOSE",
@@ -859,6 +853,12 @@
       "group": "WM_DEBUG_ANIM",
       "at": "com\/android\/server\/wm\/WindowState.java"
     },
+    "-1277068810": {
+      "message": "startBackNavigation currentTask=%s, topRunningActivity=%s, callbackInfo=%s, currentFocus=%s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_BACK_PREVIEW",
+      "at": "com\/android\/server\/wm\/BackNavigationController.java"
+    },
     "-1270731689": {
       "message": "Attempted to set replacing window on app token with no content %s",
       "level": "WARN",
@@ -1099,12 +1099,6 @@
       "group": "WM_DEBUG_STATES",
       "at": "com\/android\/server\/wm\/ActivityRecord.java"
     },
-    "-1010850753": {
-      "message": "No focused window, defaulting to top task's window",
-      "level": "WARN",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/BackNavigationController.java"
-    },
     "-1009117329": {
       "message": "isFetchingAppTransitionSpecs=true",
       "level": "VERBOSE",
@@ -3049,12 +3043,6 @@
       "group": "WM_DEBUG_REMOTE_ANIMATIONS",
       "at": "com\/android\/server\/wm\/RemoteAnimationController.java"
     },
-    "878005951": {
-      "message": "startBackNavigation task=%s, topRunningActivity=%s, callbackInfo=%s, currentFocus=%s",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/BackNavigationController.java"
-    },
     "892244061": {
       "message": "Waiting for drawn %s: removed=%b visible=%b mHasSurface=%b drawState=%d",
       "level": "INFO",
@@ -3331,12 +3319,6 @@
       "group": "WM_DEBUG_APP_TRANSITIONS",
       "at": "com\/android\/server\/wm\/DisplayContent.java"
     },
-    "1172542963": {
-      "message": "onBackNavigationDone backType=%s, task=%s, prevTaskTopActivity=%s",
-      "level": "DEBUG",
-      "group": "WM_DEBUG_BACK_PREVIEW",
-      "at": "com\/android\/server\/wm\/BackNavigationController.java"
-    },
     "1175495463": {
       "message": "ImeContainer just became organized. Reparenting under parent. imeParentSurfaceControl=%s",
       "level": "INFO",
@@ -3409,6 +3391,12 @@
       "group": "WM_DEBUG_STARTING_WINDOW",
       "at": "com\/android\/server\/wm\/ActivityRecord.java"
     },
+    "1264179654": {
+      "message": "No focused window, defaulting to top current task's window",
+      "level": "WARN",
+      "group": "WM_DEBUG_BACK_PREVIEW",
+      "at": "com\/android\/server\/wm\/BackNavigationController.java"
+    },
     "1270792394": {
       "message": "Resumed after relaunch %s",
       "level": "DEBUG",
@@ -3853,6 +3841,12 @@
       "group": "WM_ERROR",
       "at": "com\/android\/server\/wm\/WindowManagerService.java"
     },
+    "1778919449": {
+      "message": "onBackNavigationDone backType=%s, task=%s, prevActivity=%s",
+      "level": "DEBUG",
+      "group": "WM_DEBUG_BACK_PREVIEW",
+      "at": "com\/android\/server\/wm\/BackNavigationController.java"
+    },
     "1781673113": {
       "message": "onAnimationFinished(): targetRootTask=%s targetActivity=%s mRestoreTargetBehindRootTask=%s",
       "level": "DEBUG",
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index f70dc52..b37f980 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -94,12 +94,22 @@
         }
 
         int backType = BackNavigationInfo.TYPE_UNDEFINED;
+
+        // The currently visible activity (if any).
+        ActivityRecord currentActivity = null;
+
+        // The currently visible task (if any).
+        Task currentTask = null;
+
+        // The previous task we're going back to. Can be the same as currentTask, if there are
+        // multiple Activities in the Stack.
         Task prevTask = null;
-        ActivityRecord prev;
+
+        // The previous activity we're going back to. This can be either a child of currentTask
+        // if there are more than one Activity in currentTask, or a child of prevTask, if
+        // currentActivity is the last child of currentTask.
+        ActivityRecord prevActivity;
         WindowContainer<?> removedWindowContainer = null;
-        ActivityRecord activityRecord = null;
-        ActivityRecord prevTaskTopActivity = null;
-        Task task = null;
         SurfaceControl animationLeashParent = null;
         HardwareBuffer screenshotBuffer = null;
         RemoteAnimationTarget topAppTarget = null;
@@ -143,19 +153,19 @@
             }
 
             if (window == null) {
-                // We don't have any focused window, fallback ont the top task of the focused
+                // We don't have any focused window, fallback ont the top currentTask of the focused
                 // display.
                 ProtoLog.w(WM_DEBUG_BACK_PREVIEW,
-                        "No focused window, defaulting to top task's window");
-                task = wmService.mAtmService.getTopDisplayFocusedRootTask();
-                window = task.getWindow(WindowState::isFocused);
+                        "No focused window, defaulting to top current task's window");
+                currentTask = wmService.mAtmService.getTopDisplayFocusedRootTask();
+                window = currentTask.getWindow(WindowState::isFocused);
             }
 
             // Now let's find if this window has a callback from the client side.
             OnBackInvokedCallbackInfo callbackInfo = null;
             if (window != null) {
-                activityRecord = window.mActivityRecord;
-                task = window.getTask();
+                currentActivity = window.mActivityRecord;
+                currentTask = window.getTask();
                 callbackInfo = window.getOnBackInvokedCallbackInfo();
                 if (callbackInfo == null) {
                     Slog.e(TAG, "No callback registered, returning null.");
@@ -167,9 +177,9 @@
                 infoBuilder.setOnBackInvokedCallback(callbackInfo.getCallback());
             }
 
-            ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "startBackNavigation task=%s, "
+            ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "startBackNavigation currentTask=%s, "
                             + "topRunningActivity=%s, callbackInfo=%s, currentFocus=%s",
-                    task, activityRecord, callbackInfo, window);
+                    currentTask, currentActivity, callbackInfo, window);
 
             if (window == null) {
                 Slog.e(TAG, "Window is null, returning null.");
@@ -182,18 +192,18 @@
             // - The IME is opened, and we just need to close it.
             // - The home activity is the focused activity.
             if (backType == BackNavigationInfo.TYPE_CALLBACK
-                    || activityRecord == null
-                    || task == null
-                    || task.getDisplayContent().getImeContainer().isVisible()
-                    || activityRecord.isActivityTypeHome()) {
+                    || currentActivity == null
+                    || currentTask == null
+                    || currentTask.getDisplayContent().getImeContainer().isVisible()
+                    || currentActivity.isActivityTypeHome()) {
                 return infoBuilder
                         .setType(backType)
                         .build();
             }
 
             // We don't have an application callback, let's find the destination of the back gesture
-            Task finalTask = task;
-            prev = task.getActivity(
+            Task finalTask = currentTask;
+            prevActivity = currentTask.getActivity(
                     (r) -> !r.finishing && r.getTask() == finalTask && !r.isTopRunningActivity());
             if (window.getParent().getChildCount() > 1 && window.getParent().getChildAt(0)
                     != window) {
@@ -201,24 +211,24 @@
                 // activity, we won't close the activity.
                 backType = BackNavigationInfo.TYPE_DIALOG_CLOSE;
                 removedWindowContainer = window;
-            } else if (prev != null) {
-                // We have another Activity in the same task to go to
+            } else if (prevActivity != null) {
+                // We have another Activity in the same currentTask to go to
                 backType = BackNavigationInfo.TYPE_CROSS_ACTIVITY;
-                removedWindowContainer = activityRecord;
-            } else if (task.returnsToHomeRootTask()) {
+                removedWindowContainer = currentActivity;
+            } else if (currentTask.returnsToHomeRootTask()) {
                 // Our Task should bring back to home
-                removedWindowContainer = task;
+                removedWindowContainer = currentTask;
                 backType = BackNavigationInfo.TYPE_RETURN_TO_HOME;
-            } else if (activityRecord.isRootOfTask()) {
+            } else if (currentActivity.isRootOfTask()) {
                 // TODO(208789724): Create single source of truth for this, maybe in
                 //  RootWindowContainer
-                // TODO: Also check Task.shouldUpRecreateTaskLocked() for prev logic
-                prevTask = task.mRootWindowContainer.getTaskBelow(task);
-                removedWindowContainer = task;
+                // TODO: Also check Task.shouldUpRecreateTaskLocked() for prevActivity logic
+                prevTask = currentTask.mRootWindowContainer.getTaskBelow(currentTask);
+                removedWindowContainer = currentTask;
+                prevActivity = prevTask.getTopNonFinishingActivity();
                 if (prevTask.isActivityTypeHome()) {
                     backType = BackNavigationInfo.TYPE_RETURN_TO_HOME;
                 } else {
-                    prev = prevTask.getTopNonFinishingActivity();
                     backType = BackNavigationInfo.TYPE_CROSS_TASK;
                 }
             }
@@ -229,7 +239,7 @@
 
             ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "Previous Destination is Activity:%s Task:%s "
                             + "removedContainer:%s, backType=%s",
-                    prev != null ? prev.mActivityComponent : null,
+                    prevActivity != null ? prevActivity.mActivityComponent : null,
                     prevTask != null ? prevTask.getName() : null,
                     removedWindowContainer,
                     BackNavigationInfo.typeToString(backType));
@@ -241,7 +251,8 @@
                     && !removedWindowContainer.hasCommittedReparentToAnimationLeash();
 
             if (prepareAnimation) {
-                taskWindowConfiguration = task.getTaskInfo().configuration.windowConfiguration;
+                taskWindowConfiguration =
+                        currentTask.getTaskInfo().configuration.windowConfiguration;
 
                 infoBuilder.setTaskWindowConfiguration(taskWindowConfiguration);
                 // Prepare a leash to animate the current top window
@@ -254,32 +265,36 @@
                 removedWindowContainer.reparentSurfaceControl(tx, animLeash);
                 animationLeashParent = removedWindowContainer.getAnimationLeashParent();
                 topAppTarget = createRemoteAnimationTargetLocked(removedWindowContainer,
-                        activityRecord,
-                        task, animLeash);
+                        currentActivity,
+                        currentTask, animLeash);
                 infoBuilder.setDepartingAnimationTarget(topAppTarget);
             }
 
             //TODO(207481538) Remove once the infrastructure to support per-activity screenshot is
             // implemented. For now we simply have the mBackScreenshots hash map that dumbly
             // saves the screenshots.
-            if (needsScreenshot(backType) && prev != null && prev.mActivityComponent != null) {
-                screenshotBuffer = getActivitySnapshot(task, prev.mActivityComponent);
+            if (needsScreenshot(backType) && prevActivity != null
+                    && prevActivity.mActivityComponent != null) {
+                screenshotBuffer =
+                        getActivitySnapshot(currentTask, prevActivity.mActivityComponent);
             }
 
-            if (backType == BackNavigationInfo.TYPE_RETURN_TO_HOME && isAnimationEnabled()) {
-                task.mBackGestureStarted = true;
+            // Special handling for back to home animation
+            if (backType == BackNavigationInfo.TYPE_RETURN_TO_HOME && isAnimationEnabled()
+                    && prevTask != null) {
+                currentTask.mBackGestureStarted = true;
                 // Make launcher show from behind by marking its top activity as visible and
                 // launch-behind to bump its visibility for the duration of the back gesture.
-                prevTaskTopActivity = prevTask.getTopNonFinishingActivity();
-                if (prevTaskTopActivity != null) {
-                    if (!prevTaskTopActivity.mVisibleRequested) {
-                        prevTaskTopActivity.setVisibility(true);
+                prevActivity = prevTask.getTopNonFinishingActivity();
+                if (prevActivity != null) {
+                    if (!prevActivity.mVisibleRequested) {
+                        prevActivity.setVisibility(true);
                     }
-                    prevTaskTopActivity.mLaunchTaskBehind = true;
+                    prevActivity.mLaunchTaskBehind = true;
                     ProtoLog.d(WM_DEBUG_BACK_PREVIEW,
                             "Setting Activity.mLauncherTaskBehind to true. Activity=%s",
-                            prevTaskTopActivity);
-                    prevTaskTopActivity.mRootWindowContainer.ensureActivitiesVisible(
+                            prevActivity);
+                    prevActivity.mRootWindowContainer.ensureActivitiesVisible(
                             null /* starting */, 0 /* configChanges */,
                             false /* preserveWindows */);
                 }
@@ -290,7 +305,7 @@
         if (topAppTarget != null && needsScreenshot(backType) && prevTask != null
                 && screenshotBuffer == null) {
             SurfaceControl.Builder builder = new SurfaceControl.Builder()
-                    .setName("BackPreview Screenshot for " + prev)
+                    .setName("BackPreview Screenshot for " + prevActivity)
                     .setParent(animationLeashParent)
                     .setHidden(false)
                     .setBLASTLayer();
@@ -302,12 +317,12 @@
             // The Animation leash needs to be above the screenshot surface, but the animation leash
             // needs to be added before to be in the synchronized block.
             tx.setLayer(topAppTarget.leash, 1);
-            tx.apply();
+        }
 
-
-            WindowContainer<?> finalRemovedWindowContainer = removedWindowContainer;
+        WindowContainer<?> finalRemovedWindowContainer = removedWindowContainer;
+        if (finalRemovedWindowContainer != null) {
             try {
-                activityRecord.token.linkToDeath(
+                currentActivity.token.linkToDeath(
                         () -> resetSurfaces(finalRemovedWindowContainer), 0);
             } catch (RemoteException e) {
                 Slog.e(TAG, "Failed to link to death", e);
@@ -315,11 +330,16 @@
                 return null;
             }
 
-            RemoteCallback onBackNavigationDone = new RemoteCallback(
-                    result -> resetSurfaces(finalRemovedWindowContainer
-                    ));
+            int finalBackType = backType;
+            ActivityRecord finalprevActivity = prevActivity;
+            Task finalTask = currentTask;
+            RemoteCallback onBackNavigationDone = new RemoteCallback(result -> onBackNavigationDone(
+                    result, finalRemovedWindowContainer, finalBackType, finalTask,
+                    finalprevActivity));
             infoBuilder.setOnBackNavigationDone(onBackNavigationDone);
         }
+
+        tx.apply();
         return infoBuilder.build();
     }
 
@@ -348,14 +368,13 @@
     }
 
     private void onBackNavigationDone(
-            Bundle result, WindowContainer windowContainer, int backType,
-            Task task, ActivityRecord prevTaskTopActivity) {
+            Bundle result, WindowContainer<?> windowContainer, int backType,
+            Task task, ActivityRecord prevActivity) {
         SurfaceControl surfaceControl = windowContainer.getSurfaceControl();
-        boolean triggerBack = result != null
-                ? result.getBoolean(BackNavigationInfo.KEY_TRIGGER_BACK)
-                : false;
+        boolean triggerBack = result != null && result.getBoolean(
+                BackNavigationInfo.KEY_TRIGGER_BACK);
         ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "onBackNavigationDone backType=%s, "
-                + "task=%s, prevTaskTopActivity=%s", backType, task, prevTaskTopActivity);
+                + "task=%s, prevActivity=%s", backType, task, prevActivity);
 
         if (backType == BackNavigationInfo.TYPE_RETURN_TO_HOME && isAnimationEnabled()) {
             if (triggerBack) {
@@ -367,13 +386,13 @@
                     t.apply();
                 }
             }
-            if (prevTaskTopActivity != null && !triggerBack) {
+            if (prevActivity != null && !triggerBack) {
                 // Restore the launch-behind state.
-                task.mTaskSupervisor.scheduleLaunchTaskBehindComplete(prevTaskTopActivity.token);
-                prevTaskTopActivity.mLaunchTaskBehind = false;
+                task.mTaskSupervisor.scheduleLaunchTaskBehindComplete(prevActivity.token);
+                prevActivity.mLaunchTaskBehind = false;
                 ProtoLog.d(WM_DEBUG_BACK_PREVIEW,
                         "Setting Activity.mLauncherTaskBehind to false. Activity=%s",
-                        prevTaskTopActivity);
+                        prevActivity);
             }
         } else {
             task.mBackGestureStarted = false;
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
index fe59185a..49cd343 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
@@ -29,9 +29,12 @@
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.annotation.NonNull;
@@ -39,6 +42,7 @@
 import android.hardware.HardwareBuffer;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
+import android.view.SurfaceControl;
 import android.view.WindowManager;
 import android.window.BackEvent;
 import android.window.BackNavigationInfo;
@@ -77,15 +81,22 @@
 
     @Test
     public void backNavInfo_HomeWhenBackToLauncher() {
-        IOnBackInvokedCallback callback = withSystemCallback(createTopTaskWithActivity());
+        Task task = createTopTaskWithActivity();
+        IOnBackInvokedCallback callback = withSystemCallback(task);
 
-        BackNavigationInfo backNavigationInfo = startBackNavigation();
+        SurfaceControl.Transaction tx = mock(SurfaceControl.Transaction.class);
+        BackNavigationInfo backNavigationInfo = mBackNavigationController.startBackNavigation(mWm,
+                tx);
         assertWithMessage("BackNavigationInfo").that(backNavigationInfo).isNotNull();
         assertThat(backNavigationInfo.getDepartingAnimationTarget()).isNotNull();
         assertThat(backNavigationInfo.getTaskWindowConfiguration()).isNotNull();
         assertThat(backNavigationInfo.getOnBackInvokedCallback()).isEqualTo(callback);
         assertThat(typeToString(backNavigationInfo.getType()))
                 .isEqualTo(typeToString(BackNavigationInfo.TYPE_RETURN_TO_HOME));
+
+        verify(tx, atLeastOnce()).apply();
+        verify(tx, times(1)).reparent(any(),
+                eq(backNavigationInfo.getDepartingAnimationTarget().leash));
     }
 
     @Test