Properly clean up screenshot of recents animation upon cancelation

Bug: 122593881
Test: Manual
Change-Id: Ia9ce1ede08309a0898c622bcd07a9e076443d98a
(cherry picked from commit 2bd3a0225b373c21774e7f0806ff49853d410042)
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsAnimationWrapper.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsAnimationWrapper.java
index f0bc223..0924f38 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsAnimationWrapper.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/RecentsAnimationWrapper.java
@@ -173,6 +173,12 @@
         return true;
     }
 
+    public void setCancelWithDeferredScreenshot(boolean deferredWithScreenshot) {
+        if (targetSet != null) {
+            targetSet.controller.setCancelWithDeferredScreenshot(deferredWithScreenshot);
+        }
+    }
+
     public SwipeAnimationTargetSet getController() {
         return targetSet;
     }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
index eb1e7b4..e8d4c19 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -1136,6 +1136,7 @@
             mLauncherTransitionController = null;
         }
         mActivityControlHelper.onSwipeUpComplete(mActivity);
+        mRecentsAnimationWrapper.setCancelWithDeferredScreenshot(true);
 
         // Animate the first icon.
         mRecentsView.animateUpRunningTaskIconScale(mLiveTileOverlay.cancelIconAnimation());
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java
index 62f2183..94e704a 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java
@@ -38,9 +38,16 @@
  */
 public class RecentsAnimationListenerSet implements RecentsAnimationListener {
 
+    // The actual app surface is replaced by a screenshot upon recents animation cancelation when
+    // deferredWithScreenshot is true. Launcher takes the responsibility to clean up this screenshot
+    // after app transition is finished. This delay is introduced to cover the app transition
+    // period of time.
+    private final int TRANSITION_DELAY = 100;
+
     private final Set<SwipeAnimationListener> mListeners = new ArraySet<>();
     private final boolean mShouldMinimizeSplitScreen;
     private final Consumer<SwipeAnimationTargetSet> mOnFinishListener;
+    private RecentsAnimationControllerCompat mController;
 
     public RecentsAnimationListenerSet(boolean shouldMinimizeSplitScreen,
             Consumer<SwipeAnimationTargetSet> onFinishListener) {
@@ -64,6 +71,7 @@
     public final void onAnimationStart(RecentsAnimationControllerCompat controller,
             RemoteAnimationTargetCompat[] targets, Rect homeContentInsets,
             Rect minimizedHomeBounds) {
+        mController = controller;
         SwipeAnimationTargetSet targetSet = new SwipeAnimationTargetSet(controller, targets,
                 homeContentInsets, minimizedHomeBounds, mShouldMinimizeSplitScreen,
                 mOnFinishListener);
@@ -75,12 +83,17 @@
     }
 
     @Override
-    public final void onAnimationCanceled() {
+    public final void onAnimationCanceled(boolean deferredWithScreenshot) {
         Utilities.postAsyncCallback(MAIN_THREAD_EXECUTOR.getHandler(), () -> {
             for (SwipeAnimationListener listener : getListeners()) {
                 listener.onRecentsAnimationCanceled();
             }
         });
+        // TODO: handle the transition better instead of simply using a transition delay.
+        if (deferredWithScreenshot) {
+            MAIN_THREAD_EXECUTOR.getHandler().postDelayed(() -> mController.cleanupScreenshot(),
+                    TRANSITION_DELAY);
+        }
     }
 
     private SwipeAnimationListener[] getListeners() {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
index 8f2a2d7..8da1d2b 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -74,12 +74,7 @@
 
     @Override
     public void startHome() {
-        if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
-            takeScreenshotAndFinishRecentsAnimation(true,
-                    () -> mActivity.getStateManager().goToState(NORMAL));
-        } else {
-            mActivity.getStateManager().goToState(NORMAL);
-        }
+        mActivity.getStateManager().goToState(NORMAL);
     }
 
     @Override
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
index 37febf9a..a02df62 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
@@ -96,7 +96,6 @@
 import com.android.quickstep.TaskThumbnailCache;
 import com.android.quickstep.TaskUtils;
 import com.android.quickstep.util.ClipAnimationHelper;
-import com.android.quickstep.util.SwipeAnimationTargetSet;
 import com.android.quickstep.util.TaskViewDrawable;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -106,7 +105,6 @@
 import com.android.systemui.shared.system.PackageManagerWrapper;
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
 import com.android.systemui.shared.system.TaskStackChangeListener;
-import com.android.systemui.shared.system.WindowCallbacksCompat;
 
 import java.util.ArrayList;
 import java.util.function.Consumer;
@@ -1608,50 +1606,4 @@
 
         mRecentsAnimationWrapper.finish(toRecents, onFinishComplete);
     }
-
-    public void takeScreenshotAndFinishRecentsAnimation(boolean toRecents,
-            Runnable onFinishComplete) {
-        if (mRecentsAnimationWrapper == null || getRunningTaskView() == null) {
-            if (onFinishComplete != null) {
-                onFinishComplete.run();
-            }
-            return;
-        }
-
-        SwipeAnimationTargetSet controller = mRecentsAnimationWrapper.getController();
-        if (controller != null) {
-            // Update the screenshot of the task
-            ThumbnailData taskSnapshot = controller.screenshotTask(mRunningTaskId);
-            TaskView taskView = updateThumbnail(mRunningTaskId, taskSnapshot);
-            if (taskView != null) {
-                taskView.setShowScreenshot(true);
-                // Defer finishing the animation until the next launcher frame with the
-                // new thumbnail
-                new WindowCallbacksCompat(taskView) {
-
-                    // The number of frames to defer until we actually finish the animation
-                    private int mDeferFrameCount = 2;
-
-                    @Override
-                    public void onPostDraw(Canvas canvas) {
-                        if (mDeferFrameCount > 0) {
-                            mDeferFrameCount--;
-                            // Workaround, detach and reattach to invalidate the root node for
-                            // another draw
-                            detach();
-                            attach();
-                            taskView.invalidate();
-                            return;
-                        }
-
-                        detach();
-                        mRecentsAnimationWrapper.finish(toRecents, () -> {
-                            onFinishComplete.run();
-                            mRunningTaskId = -1;
-                        });
-                    }
-                }.attach();
-            }
-        }
-    }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java
index 682152e..d15a392 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskMenuView.java
@@ -208,13 +208,7 @@
                 R.layout.task_view_menu_option, this, false);
         menuOption.setIconAndLabelFor(
                 menuOptionView.findViewById(R.id.icon), menuOptionView.findViewById(R.id.text));
-        if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
-            menuOptionView.setOnClickListener(
-                    view -> mTaskView.getRecentsView().takeScreenshotAndFinishRecentsAnimation(true,
-                            () -> onClickListener.onClick(view)));
-        } else {
-            menuOptionView.setOnClickListener(onClickListener);
-        }
+        menuOptionView.setOnClickListener(onClickListener);
         mOptionLayout.addView(menuOptionView);
     }
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
index 98495db..38aaac5 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
@@ -261,11 +261,10 @@
             Handler resultCallbackHandler) {
         if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
             if (isRunningTask()) {
-                getRecentsView().finishRecentsAnimation(false,
+                getRecentsView().finishRecentsAnimation(false /* toRecents */,
                         () -> resultCallbackHandler.post(() -> resultCallback.accept(true)));
             } else {
-                getRecentsView().takeScreenshotAndFinishRecentsAnimation(true,
-                        () -> launchTaskInternal(animate, resultCallback, resultCallbackHandler));
+                launchTaskInternal(animate, resultCallback, resultCallbackHandler);
             }
         } else {
             launchTaskInternal(animate, resultCallback, resultCallbackHandler);