Fix flash when switching from app window to snapshot.

- If there is any delay in launcher draw, then we will finish the
  transition and hide the app window before launcher has had a chance to
  draw the task view with the updated snapshot.

Change-Id: I81c5b467f32520fd7474bde0e89819336292e84f
diff --git a/quickstep/src/com/android/quickstep/RecentsView.java b/quickstep/src/com/android/quickstep/RecentsView.java
index c400ffd..c8122d2 100644
--- a/quickstep/src/com/android/quickstep/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/RecentsView.java
@@ -132,15 +132,16 @@
         mScrollState.isRtl = mIsRtl;
     }
 
-    public void updateThumbnail(int taskId, ThumbnailData thumbnailData) {
+    public TaskView updateThumbnail(int taskId, ThumbnailData thumbnailData) {
         for (int i = mFirstTaskIndex; i < getChildCount(); i++) {
             final TaskView taskView = (TaskView) getChildAt(i);
             if (taskView.getTask().key.id == taskId) {
                 taskView.onTaskDataLoaded(taskView.getTask(), thumbnailData);
                 taskView.setAlpha(1);
-                return;
+                return taskView;
             }
         }
+        return null;
     }
 
     private void setupLayoutTransition() {
diff --git a/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java b/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java
index a6192c2..012afb2 100644
--- a/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -46,7 +46,9 @@
 import android.support.annotation.WorkerThread;
 import android.util.Log;
 import android.view.View;
+import android.view.ViewTreeObserver;
 import android.view.ViewTreeObserver.OnDrawListener;
+import android.view.ViewTreeObserver.OnPreDrawListener;
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.DeviceProfile;
@@ -67,6 +69,7 @@
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action.Direction;
 import com.android.launcher3.userevent.nano.LauncherLogProto.ContainerType;
 import com.android.launcher3.util.TraceHelper;
+import com.android.launcher3.util.ViewOnDrawExecutor;
 import com.android.quickstep.TouchConsumer.InteractionType;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.recents.utilities.RectFEvaluator;
@@ -214,22 +217,24 @@
         mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_GESTURE_STARTED,
                 this::notifyGestureStarted);
         mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_LAUNCHER_STARTED
-                | STATE_GESTURE_CANCELLED, this::resetStateForAnimationCancel);
+                        | STATE_GESTURE_CANCELLED,
+                this::resetStateForAnimationCancel);
 
-        mStateCallback.addCallback(STATE_SCALED_CONTROLLER_APP | STATE_APP_CONTROLLER_RECEIVED,
+        mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_APP_CONTROLLER_RECEIVED
+                        | STATE_SCALED_CONTROLLER_APP,
                 this::resumeLastTask);
-        mStateCallback.addCallback(STATE_SCALED_CONTROLLER_RECENTS
-                | STATE_ACTIVITY_MULTIPLIER_COMPLETE | STATE_APP_CONTROLLER_RECEIVED,
+        mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_APP_CONTROLLER_RECEIVED
+                        | STATE_ACTIVITY_MULTIPLIER_COMPLETE
+                        | STATE_SCALED_CONTROLLER_RECENTS,
                 this::switchToScreenshot);
-
-        mStateCallback.addCallback(STATE_SCALED_CONTROLLER_RECENTS
-                        | STATE_ACTIVITY_MULTIPLIER_COMPLETE,
+        mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_APP_CONTROLLER_RECEIVED
+                        | STATE_ACTIVITY_MULTIPLIER_COMPLETE
+                        | STATE_SCALED_CONTROLLER_RECENTS
+                        | STATE_SWITCH_TO_SCREENSHOT_COMPLETE,
                 this::setupLauncherUiAfterSwipeUpAnimation);
 
         mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_SCALED_CONTROLLER_APP,
                 this::reset);
-        mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_SCALED_CONTROLLER_RECENTS,
-                this::reset);
 
         mStateCallback.addCallback(STATE_HANDLER_INVALIDATED, this::invalidateHandler);
         mStateCallback.addCallback(STATE_LAUNCHER_PRESENT | STATE_HANDLER_INVALIDATED,
@@ -718,22 +723,47 @@
     }
 
     private void switchToScreenshot() {
+        boolean finishTransitionPosted = false;
+        final Runnable finishTransitionRunnable = () -> {
+            synchronized (mRecentsAnimationWrapper) {
+                mRecentsAnimationWrapper.finish(true /* toHome */,
+                        () -> setStateOnUiThread(STATE_SWITCH_TO_SCREENSHOT_COMPLETE));
+            }
+        };
         synchronized (mRecentsAnimationWrapper) {
             if (mRecentsAnimationWrapper.controller != null) {
                 TransactionCompat transaction = new TransactionCompat();
                 for (RemoteAnimationTargetCompat app : mRecentsAnimationWrapper.targets) {
                     if (app.mode == MODE_CLOSING) {
                         // Update the screenshot of the task
-                        final ThumbnailData thumbnail =
+                        ThumbnailData thumbnail =
                                 mRecentsAnimationWrapper.controller.screenshotTask(app.taskId);
-                        mRecentsView.updateThumbnail(app.taskId, thumbnail);
+                        TaskView taskView = mRecentsView.updateThumbnail(app.taskId, thumbnail);
+                        if (taskView != null) {
+                            // Defer finishing the animation until the next launcher frame with the
+                            // new thumbnail
+                            ViewOnDrawExecutor executor = new ViewOnDrawExecutor(mMainExecutor) {
+                                @Override
+                                public void onViewDetachedFromWindow(View v) {
+                                    if (!isCompleted()) {
+                                        runAllTasks();
+                                    }
+                                }
+                            };
+                            executor.attachTo(mLauncher, taskView,
+                                    false /* waitForLoadAnimation */);
+                            executor.execute(finishTransitionRunnable);
+                            finishTransitionPosted = true;
+                        }
                     }
                 }
                 transaction.apply();
             }
         }
-        mRecentsAnimationWrapper.finish(true /* toHome */,
-                () -> setStateOnUiThread(STATE_SWITCH_TO_SCREENSHOT_COMPLETE));
+        if (!finishTransitionPosted) {
+            // If we haven't posted the transition end runnable, run it now
+            finishTransitionRunnable.run();
+        }
         doLogGesture(true /* toLauncher */);
     }
 
@@ -743,6 +773,8 @@
 
         // Animate the first icon.
         mRecentsView.setFirstTaskIconScaledDown(false /* isScaledDown */, true /* animate */);
+
+        reset();
     }
 
     public void onQuickScrubEnd() {
diff --git a/src/com/android/launcher3/util/ViewOnDrawExecutor.java b/src/com/android/launcher3/util/ViewOnDrawExecutor.java
index e5c1dd1..34bdcc9 100644
--- a/src/com/android/launcher3/util/ViewOnDrawExecutor.java
+++ b/src/com/android/launcher3/util/ViewOnDrawExecutor.java
@@ -49,9 +49,16 @@
     }
 
     public void attachTo(Launcher launcher) {
+        attachTo(launcher, launcher.getWorkspace(), true /* waitForLoadAnimation */);
+    }
+
+    public void attachTo(Launcher launcher, View attachedView, boolean waitForLoadAnimation) {
         mLauncher = launcher;
-        mAttachedView = launcher.getWorkspace();
+        mAttachedView = attachedView;
         mAttachedView.addOnAttachStateChangeListener(this);
+        if (!waitForLoadAnimation) {
+            mLoadAnimationCompleted = true;
+        }
 
         attachObserver();
     }
@@ -74,7 +81,7 @@
     }
 
     @Override
-    public void onViewDetachedFromWindow(View v) { }
+    public void onViewDetachedFromWindow(View v) {}
 
     @Override
     public void onDraw() {
@@ -101,10 +108,7 @@
         // Post the pending tasks after both onDraw and onLoadAnimationCompleted have been called.
         if (mLoadAnimationCompleted && mFirstDrawCompleted && !mCompleted) {
             mIsExecuting = true;
-            for (final Runnable r : mTasks) {
-                mExecutor.execute(r);
-            }
-            markCompleted();
+            runAllTasks();
         }
     }
 
@@ -121,4 +125,15 @@
         }
         LauncherModel.setWorkerPriority(Process.THREAD_PRIORITY_DEFAULT);
     }
+
+    protected boolean isCompleted() {
+        return mCompleted;
+    }
+
+    protected void runAllTasks() {
+        for (final Runnable r : mTasks) {
+            mExecutor.execute(r);
+        }
+        markCompleted();
+    }
 }