Simplify UI and animation choreography a bit

A number of changes to improve the robustness of the
startup/reset/dismiss process of the screenshot UI.

Brief summary:

Use visibility=invisible always for preview image
Wait for pre-draw to start animation
Apply initial config to "interesting changes" tracker
Rename: mOnCompleteAction -> mOnDismissedAction
Rease more resources earlier at shutdown
Remove some uneccessary main thread deferrals
resetScreenshotView -> finishDismiss
Resolve some cases where reset is called multiple times on exit
Update window token usage and apply to scroll capture at window attach
Simplified animation parameters, coordinate systems, etc.

Bug: 174509231
Bug: 174682382
Test: manual; take a screenshot, dismiss, share, edit, swipe away
Test: manual; open overview, tap screenshot

Change-Id: I3b84da3592970ac5c2a5eb34326eabaef003993f
Merged-In: I3b84da3592970ac5c2a5eb34326eabaef003993f

5 files changed, 246 insertions, 163 deletions

diff --git a/packages/SystemUI/res/layout/global_screenshot_preview.xml b/packages/SystemUI/res/layout/global_screenshot_preview.xml
index 5262407ffef9..100213bdfd8c 100644
--- a/packages/SystemUI/res/layout/global_screenshot_preview.xml
+++ b/packages/SystemUI/res/layout/global_screenshot_preview.xml
@@ -25,7 +25,7 @@
     android:layout_marginBottom="@dimen/screenshot_offset_y"
     android:scaleType="fitEnd"
     android:elevation="@dimen/screenshot_preview_elevation"
-    android:visibility="gone"
+    android:visibility="invisible"
     android:background="@drawable/screenshot_rounded_corners"
     android:adjustViewBounds="true"
     android:contentDescription="@string/screenshot_edit_label"
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index 395729ce2c94..4431b6974b7b 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -66,6 +66,7 @@ import android.view.LayoutInflater;
 import android.view.MotionEvent;
 import android.view.SurfaceControl;
 import android.view.View;
+import android.view.ViewTreeObserver;
 import android.view.Window;
 import android.view.WindowInsets;
 import android.view.WindowManager;
@@ -79,13 +80,15 @@ import com.android.internal.logging.UiEventLogger;
 import com.android.internal.policy.PhoneWindow;
 import com.android.settingslib.applications.InterestingConfigChanges;
 import com.android.systemui.R;
-import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.screenshot.ScreenshotController.SavedImageData.ActionTransition;
+import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback;
 import com.android.systemui.util.DeviceConfigProxy;
 
 import java.util.List;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
@@ -172,7 +175,7 @@ public class ScreenshotController {
     private final UiEventLogger mUiEventLogger;
     private final ImageExporter mImageExporter;
     private final Executor mMainExecutor;
-    private final Executor mBgExecutor;
+    private final ExecutorService mBgExecutor;
 
     private final WindowManager mWindowManager;
     private final WindowManager.LayoutParams mWindowLayoutParams;
@@ -181,7 +184,6 @@ public class ScreenshotController {
     private final ScrollCaptureClient mScrollCaptureClient;
     private final DeviceConfigProxy mConfigProxy;
     private final PhoneWindow mWindow;
-    private final View mDecorView;
     private final DisplayManager mDisplayManager;
 
     private ScreenshotView mScreenshotView;
@@ -189,8 +191,7 @@ public class ScreenshotController {
     private SaveImageInBackgroundTask mSaveInBgTask;
 
     private Animator mScreenshotAnimation;
-
-    private Runnable mOnCompleteRunnable;
+    private RequestCallback mCurrentRequestCallback;
 
     private final Handler mScreenshotHandler = new Handler(Looper.getMainLooper()) {
         @Override
@@ -227,15 +228,14 @@ public class ScreenshotController {
             UiEventLogger uiEventLogger,
             DeviceConfigProxy configProxy,
             ImageExporter imageExporter,
-            @Main Executor mainExecutor,
-            @Background Executor bgExecutor) {
+            @Main Executor mainExecutor) {
         mScreenshotSmartActions = screenshotSmartActions;
         mNotificationsController = screenshotNotificationsController;
         mScrollCaptureClient = scrollCaptureClient;
         mUiEventLogger = uiEventLogger;
         mImageExporter = imageExporter;
         mMainExecutor = mainExecutor;
-        mBgExecutor = bgExecutor;
+        mBgExecutor = Executors.newSingleThreadExecutor();
 
         mDisplayManager = requireNonNull(context.getSystemService(DisplayManager.class));
         final Context displayContext = context.createDisplayContext(getDefaultDisplay());
@@ -266,8 +266,8 @@ public class ScreenshotController {
         mWindow.requestFeature(Window.FEATURE_NO_TITLE);
         mWindow.requestFeature(Window.FEATURE_ACTIVITY_TRANSITIONS);
         mWindow.setBackgroundDrawableResource(android.R.color.transparent);
-        mDecorView = mWindow.getDecorView();
 
+        mConfigChanges.applyNewConfig(context.getResources());
         reloadAssets();
 
         // Setup the Camera shutter sound
@@ -275,9 +275,8 @@ public class ScreenshotController {
         mCameraSound.load(MediaActionSound.SHUTTER_CLICK);
     }
 
-    void takeScreenshotFullscreen(Consumer<Uri> finisher, Runnable onComplete) {
-        mOnCompleteRunnable = onComplete;
-
+    void takeScreenshotFullscreen(Consumer<Uri> finisher, RequestCallback requestCallback) {
+        mCurrentRequestCallback = requestCallback;
         DisplayMetrics displayMetrics = new DisplayMetrics();
         getDefaultDisplay().getRealMetrics(displayMetrics);
         takeScreenshotInternal(
@@ -287,16 +286,14 @@ public class ScreenshotController {
 
     void handleImageAsScreenshot(Bitmap screenshot, Rect screenshotScreenBounds,
             Insets visibleInsets, int taskId, int userId, ComponentName topComponent,
-            Consumer<Uri> finisher, Runnable onComplete) {
+            Consumer<Uri> finisher, RequestCallback requestCallback) {
         // TODO: use task Id, userId, topComponent for smart handler
-        mOnCompleteRunnable = onComplete;
 
         if (screenshot == null) {
             Log.e(TAG, "Got null bitmap from screenshot message");
             mNotificationsController.notifyScreenshotError(
                     R.string.screenshot_failed_to_capture_text);
-            finisher.accept(null);
-            mOnCompleteRunnable.run();
+            requestCallback.reportError();
             return;
         }
 
@@ -306,17 +303,19 @@ public class ScreenshotController {
             visibleInsets = Insets.NONE;
             screenshotScreenBounds.set(0, 0, screenshot.getWidth(), screenshot.getHeight());
         }
+        mCurrentRequestCallback = requestCallback;
         saveScreenshot(screenshot, finisher, screenshotScreenBounds, visibleInsets, showFlash);
     }
 
     /**
      * Displays a screenshot selector
      */
-    void takeScreenshotPartial(final Consumer<Uri> finisher, Runnable onComplete) {
-        dismissScreenshot(true);
-        mOnCompleteRunnable = onComplete;
+    void takeScreenshotPartial(final Consumer<Uri> finisher, RequestCallback requestCallback) {
+        mScreenshotView.reset();
+        mCurrentRequestCallback = requestCallback;
 
-        mWindowManager.addView(mScreenshotView, mWindowLayoutParams);
+        attachWindow();
+        mWindow.setContentView(mScreenshotView);
 
         mScreenshotView.takePartialScreenshot(
                 rect -> takeScreenshotInternal(finisher, rect));
@@ -337,9 +336,9 @@ public class ScreenshotController {
             }
             return;
         }
-        mScreenshotHandler.removeMessages(MESSAGE_CORNER_TIMEOUT);
+        cancelTimeout();
         if (immediate) {
-            resetScreenshotView();
+            finishDismiss();
         } else {
             mScreenshotView.animateDismissal();
         }
@@ -354,6 +353,8 @@ public class ScreenshotController {
      */
     void releaseContext() {
         mContext.release();
+        mCameraSound.release();
+        mBgExecutor.shutdownNow();
     }
 
     /**
@@ -363,12 +364,8 @@ public class ScreenshotController {
         if (DEBUG_UI) {
             Log.d(TAG, "reloadAssets()");
         }
-        boolean wasAttached = mDecorView.isAttachedToWindow();
-        if (wasAttached) {
-            if (DEBUG_WINDOW) {
-                Log.d(TAG, "Removing screenshot window");
-            }
-            mWindowManager.removeView(mDecorView);
+        if (mScreenshotView != null && mScreenshotView.isAttachedToWindow()) {
+            mWindow.clearContentView(); // Is there a simpler way to say "remove screenshotView?"
         }
 
         // respect the display cutout in landscape (since we'd otherwise overlap) but not portrait
@@ -376,12 +373,6 @@ public class ScreenshotController {
         mWindowLayoutParams.setFitInsetsTypes(
                 orientation == ORIENTATION_PORTRAIT ? 0 : WindowInsets.Type.displayCutout());
 
-        // ignore system bar insets for the purpose of window layout
-        mDecorView.setOnApplyWindowInsetsListener((v, insets) -> v.onApplyWindowInsets(
-                new WindowInsets.Builder(insets)
-                        .setInsets(WindowInsets.Type.all(), Insets.NONE)
-                        .build()));
-
         // Inflate the screenshot layout
         mScreenshotView = (ScreenshotView)
                 LayoutInflater.from(mContext).inflate(R.layout.global_screenshot, null);
@@ -393,12 +384,18 @@ public class ScreenshotController {
 
             @Override
             public void onDismiss() {
-                resetScreenshotView();
+                finishDismiss();
             }
         });
 
+        // ignore system bar insets for the purpose of window layout
+        mScreenshotView.setOnApplyWindowInsetsListener((v, insets) -> v.onApplyWindowInsets(
+                new WindowInsets.Builder(insets)
+                        .setInsets(WindowInsets.Type.all(), Insets.NONE)
+                        .build()));
+
         // TODO(159460485): Remove this when focus is handled properly in the system
-        mDecorView.setOnTouchListener((v, event) -> {
+        mScreenshotView.setOnTouchListener((v, event) -> {
             if (event.getActionMasked() == MotionEvent.ACTION_OUTSIDE) {
                 if (DEBUG_INPUT) {
                     Log.d(TAG, "onTouch: ACTION_OUTSIDE");
@@ -420,8 +417,10 @@ public class ScreenshotController {
             return false;
         });
 
-        // view is added to window manager in startAnimation
-        mWindow.setContentView(mScreenshotView, mWindowLayoutParams);
+        if (DEBUG_WINDOW) {
+            Log.d(TAG, "adding OnComputeInternalInsetsListener");
+        }
+        mScreenshotView.getViewTreeObserver().addOnComputeInternalInsetsListener(mScreenshotView);
     }
 
     /**
@@ -458,14 +457,9 @@ public class ScreenshotController {
             Log.e(TAG, "takeScreenshotInternal: Screenshot bitmap was null");
             mNotificationsController.notifyScreenshotError(
                     R.string.screenshot_failed_to_capture_text);
-            if (DEBUG_CALLBACK) {
-                Log.d(TAG, "Supplying null to Consumer<Uri>");
-            }
-            finisher.accept(null);
-            if (DEBUG_CALLBACK) {
-                Log.d(TAG, "Calling mOnCompleteRunnable.run()");
+            if (mCurrentRequestCallback != null) {
+                mCurrentRequestCallback.reportError();
             }
-            mOnCompleteRunnable.run();
             return;
         }
 
@@ -482,6 +476,13 @@ public class ScreenshotController {
             mAccessibilityManager.sendAccessibilityEvent(event);
         }
 
+        if (mConfigChanges.applyNewConfig(mContext.getResources())) {
+            if (DEBUG_UI) {
+                Log.d(TAG, "saveScreenshot: reloading assets");
+            }
+            reloadAssets();
+        }
+
         if (mScreenshotView.isAttachedToWindow()) {
             // if we didn't already dismiss for another reason
             if (!mScreenshotView.isDismissing()) {
@@ -508,18 +509,101 @@ public class ScreenshotController {
         mScreenBitmap.setHasAlpha(false);
         mScreenBitmap.prepareToDraw();
 
-        if (mConfigChanges.applyNewConfig(mContext.getResources())) {
-            if (DEBUG_UI) {
-                Log.d(TAG, "saveScreenshot: reloading assets");
-            }
-            reloadAssets();
-        }
+        saveScreenshotInWorkerThread(finisher, this::showUiOnActionsReady);
 
         // The window is focusable by default
         setWindowFocusable(true);
 
-        // Start the post-screenshot animation
-        startAnimation(finisher, screenRect, screenInsets, showFlash);
+        if (mConfigProxy.getBoolean(DeviceConfig.NAMESPACE_SYSTEMUI,
+                SystemUiDeviceConfigFlags.SCREENSHOT_SCROLLING_ENABLED, false)) {
+            View decorView = mWindow.getDecorView();
+
+            // Wait until this window is attached to request because it is
+            // the reference used to locate the target window (below).
+            withWindowAttached(() -> {
+                mScrollCaptureClient.setHostWindowToken(decorView.getWindowToken());
+                mScrollCaptureClient.request(DEFAULT_DISPLAY,
+                        /* onConnection */
+                        (connection) -> mScreenshotHandler.post(() ->
+                                mScreenshotView.showScrollChip(() ->
+                                        /* onClick */
+                                        runScrollCapture(connection))));
+            });
+        }
+
+
+        attachWindow();
+        if (DEBUG_WINDOW) {
+            Log.d(TAG, "setContentView: " + mScreenshotView);
+        }
+        mScreenshotView.getViewTreeObserver().addOnPreDrawListener(
+                new ViewTreeObserver.OnPreDrawListener() {
+                    @Override
+                    public boolean onPreDraw() {
+                        if (DEBUG_WINDOW) {
+                            Log.d(TAG, "onPreDraw: startAnimation");
+                        }
+                        mScreenshotView.getViewTreeObserver().removeOnPreDrawListener(this);
+                        startAnimation(screenRect, showFlash);
+                        return true;
+                    }
+                });
+        mScreenshotView.setScreenshot(mScreenBitmap, screenInsets);
+        setContentView(mScreenshotView);
+        cancelTimeout(); // restarted after animation
+    }
+
+    private void withWindowAttached(Runnable action) {
+        View decorView = mWindow.getDecorView();
+        if (decorView.isAttachedToWindow()) {
+            action.run();
+        } else {
+            decorView.getViewTreeObserver().addOnWindowAttachListener(
+                    new ViewTreeObserver.OnWindowAttachListener() {
+                        @Override
+                        public void onWindowAttached() {
+                            decorView.getViewTreeObserver().removeOnWindowAttachListener(this);
+                            action.run();
+                        }
+
+                        @Override
+                        public void onWindowDetached() { }
+                    });
+
+        }
+    }
+
+    private void setContentView(View contentView) {
+        mWindow.setContentView(contentView);
+    }
+
+    private void attachWindow() {
+        View decorView = mWindow.getDecorView();
+        if (decorView.isAttachedToWindow()) {
+            return;
+        }
+        if (DEBUG_WINDOW) {
+            Log.d(TAG, "attachWindow");
+        }
+        mWindowManager.addView(decorView, mWindowLayoutParams);
+        decorView.requestApplyInsets();
+    }
+
+    void removeWindow() {
+        final View decorView = mWindow.peekDecorView();
+        if (decorView != null && decorView.isAttachedToWindow()) {
+            if (DEBUG_WINDOW) {
+                Log.d(TAG, "Removing screenshot window");
+            }
+            mWindowManager.removeViewImmediate(decorView);
+        }
+    }
+
+    private void runScrollCapture(ScrollCaptureClient.Connection connection) {
+        cancelTimeout();
+        ScrollCaptureController controller = new ScrollCaptureController(mContext, connection,
+                mMainExecutor, mBgExecutor, mImageExporter);
+        controller.start(/* onDismiss */ () -> dismissScreenshot(false));
     }
 
     /**
@@ -528,11 +612,11 @@ public class ScreenshotController {
      */
     private void saveScreenshotAndToast(Consumer<Uri> finisher) {
         // Play the shutter sound to notify that we've taken a screenshot
-        mScreenshotHandler.post(() -> {
-            mCameraSound.play(MediaActionSound.SHUTTER_CLICK);
-        });
+        mCameraSound.play(MediaActionSound.SHUTTER_CLICK);
 
-        saveScreenshotInWorkerThread(finisher, imageData -> {
+        saveScreenshotInWorkerThread(
+                /* onComplete */ finisher,
+                /* actionsReadyListener */ imageData -> {
             if (DEBUG_CALLBACK) {
                 Log.d(TAG, "returning URI to finisher (Consumer<URI>): " + imageData.uri);
             }
@@ -550,88 +634,37 @@ public class ScreenshotController {
     }
 
     /**
-     * If scrolling is enabled, check whether the current view is scrollable and if so, show the
-     * scroll chip.
-     */
-    private void maybeRequestScrollCapture() {
-        if (mConfigProxy.getBoolean(DeviceConfig.NAMESPACE_SYSTEMUI,
-                SystemUiDeviceConfigFlags.SCREENSHOT_SCROLLING_ENABLED, false)) {
-            mScrollCaptureClient.setHostWindowToken(mDecorView.getWindowToken());
-            mScrollCaptureClient.request(DEFAULT_DISPLAY, (connection) ->
-                    mScreenshotView.showScrollChip(() -> {
-                        ScrollCaptureController controller = new ScrollCaptureController(
-                                mContext, connection, mMainExecutor, mBgExecutor, mImageExporter);
-                        controller.run(() -> mScreenshotHandler.post(
-                                () -> dismissScreenshot(false)));
-                    }));
-        }
-    }
-
-    /**
      * Starts the animation after taking the screenshot
      */
-    private void startAnimation(final Consumer<Uri> finisher, Rect screenRect, Insets screenInsets,
-            boolean showFlash) {
-        mScreenshotHandler.removeMessages(MESSAGE_CORNER_TIMEOUT);
-        mScreenshotHandler.post(() -> {
-            if (!mDecorView.isAttachedToWindow()) {
-                if (DEBUG_WINDOW) {
-                    Log.d(TAG, "Adding screenshot window");
-                }
-                mDecorView.addOnAttachStateChangeListener(new View.OnAttachStateChangeListener() {
-                    @Override
-                    public void onViewAttachedToWindow(View v) {
-                        mDecorView.removeOnAttachStateChangeListener(this);
-                        maybeRequestScrollCapture();
-                    }
-
-                    @Override
-                    public void onViewDetachedFromWindow(View v) {
-                    }
-                });
-                mWindowManager.addView(mWindow.getDecorView(), mWindowLayoutParams);
-            } else {
-                maybeRequestScrollCapture();
-            }
-
-            mScreenshotView.prepareForAnimation(mScreenBitmap, screenInsets);
-
-            mScreenshotHandler.post(() -> {
-                if (DEBUG_WINDOW) {
-                    Log.d(TAG, "adding OnComputeInternalInsetsListener");
-                }
-                mScreenshotView.getViewTreeObserver().addOnComputeInternalInsetsListener(
-                        mScreenshotView);
-
-                mScreenshotAnimation =
-                        mScreenshotView.createScreenshotDropInAnimation(screenRect, showFlash);
+    private void startAnimation(Rect screenRect, boolean showFlash) {
+        if (mScreenshotAnimation != null && mScreenshotAnimation.isRunning()) {
+            mScreenshotAnimation.cancel();
+        }
 
-                saveScreenshotInWorkerThread(finisher, this::showUiOnActionsReady);
+        mScreenshotAnimation =
+                mScreenshotView.createScreenshotDropInAnimation(screenRect, showFlash);
 
-                // Play the shutter sound to notify that we've taken a screenshot
-                mCameraSound.play(MediaActionSound.SHUTTER_CLICK);
+        // Play the shutter sound to notify that we've taken a screenshot
+        mCameraSound.play(MediaActionSound.SHUTTER_CLICK);
 
-                if (DEBUG_ANIM) {
-                    Log.d(TAG, "starting post-screenshot animation");
-                }
-                mScreenshotAnimation.start();
-            });
-        });
+        if (DEBUG_ANIM) {
+            Log.d(TAG, "starting post-screenshot animation");
+        }
+        mScreenshotAnimation.start();
     }
 
     /** Reset screenshot view and then call onCompleteRunnable */
-    private void resetScreenshotView() {
+    private void finishDismiss() {
         if (DEBUG_UI) {
-            Log.d(TAG, "resetScreenshotView");
+            Log.d(TAG, "finishDismiss");
         }
-        if (mScreenshotView.isAttachedToWindow()) {
-            if (DEBUG_WINDOW) {
-                Log.d(TAG, "Removing screenshot window");
-            }
-            mWindowManager.removeViewImmediate(mDecorView);
-        }
-        mOnCompleteRunnable.run();
+        cancelTimeout();
+        removeWindow();
         mScreenshotView.reset();
+        if (mCurrentRequestCallback != null) {
+            mCurrentRequestCallback.onFinish();
+            mCurrentRequestCallback = null;
+        }
     }
 
     /**
@@ -655,8 +688,12 @@ public class ScreenshotController {
         mSaveInBgTask.execute();
     }
 
-    private void resetTimeout() {
+    private void cancelTimeout() {
         mScreenshotHandler.removeMessages(MESSAGE_CORNER_TIMEOUT);
+    }
+
+    private void resetTimeout() {
+        cancelTimeout();
 
         AccessibilityManager accessibilityManager = (AccessibilityManager)
                 mContext.getSystemService(Context.ACCESSIBILITY_SERVICE);
@@ -715,7 +752,7 @@ public class ScreenshotController {
 
                 @Override
                 public void hideSharedElements() {
-                    resetScreenshotView();
+                    finishDismiss();
                 }
 
                 @Override
@@ -763,13 +800,21 @@ public class ScreenshotController {
         if (DEBUG_WINDOW) {
             Log.d(TAG, "setWindowFocusable: " + focusable);
         }
+        int flags = mWindowLayoutParams.flags;
         if (focusable) {
             mWindowLayoutParams.flags &= ~WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
         } else {
             mWindowLayoutParams.flags |= WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
         }
-        if (mDecorView.isAttachedToWindow()) {
-            mWindowManager.updateViewLayout(mDecorView, mWindowLayoutParams);
+        if (mWindowLayoutParams.flags == flags) {
+            if (DEBUG_WINDOW) {
+                Log.d(TAG, "setWindowFocusable: skipping, already " + focusable);
+            }
+            return;
+        }
+        final View decorView = mWindow.peekDecorView();
+        if (decorView != null && decorView.isAttachedToWindow()) {
+            mWindowManager.updateViewLayout(decorView, mWindowLayoutParams);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
index 211f5072bd1a..bf86b68893c8 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
@@ -316,21 +316,21 @@ public class ScreenshotView extends FrameLayout implements
         mScreenshotSelectorView.requestFocus();
     }
 
-    void prepareForAnimation(Bitmap bitmap, Insets screenInsets) {
+    void setScreenshot(Bitmap bitmap, Insets screenInsets) {
         mScreenshotPreview.setImageDrawable(createScreenDrawable(mResources, bitmap, screenInsets));
-        // make static preview invisible (from gone) so we can query its location on screen
-        mScreenshotPreview.setVisibility(View.INVISIBLE);
     }
 
     AnimatorSet createScreenshotDropInAnimation(Rect bounds, boolean showFlash) {
-        mScreenshotPreview.setLayerType(View.LAYER_TYPE_HARDWARE, null);
-        mScreenshotPreview.buildLayer();
+        if (DEBUG_ANIM) {
+            Log.d(TAG, "createAnim: bounds=" + bounds + " showFlash=" + showFlash);
+        }
 
         Rect previewBounds = new Rect();
         mScreenshotPreview.getBoundsOnScreen(previewBounds);
-        int[] previewLocation = new int[2];
-        mScreenshotPreview.getLocationInWindow(previewLocation);
+        Rect targetPosition = new Rect();
+        mScreenshotPreview.getHitRect(targetPosition);
 
+        // ratio of preview width, end vs. start size
         float cornerScale =
                 mCornerSizeX / (mOrientationPortrait ? bounds.width() : bounds.height());
         final float currentScale = 1 / cornerScale;
@@ -358,8 +358,13 @@ public class ScreenshotView extends FrameLayout implements
 
         // animate from the current location, to the static preview location
         final PointF startPos = new PointF(bounds.centerX(), bounds.centerY());
-        final PointF finalPos = new PointF(previewLocation[0] + previewBounds.width() / 2f,
-                previewLocation[1] + previewBounds.height() / 2f);
+        final PointF finalPos = new PointF(targetPosition.exactCenterX(),
+                targetPosition.exactCenterY());
+
+        if (DEBUG_ANIM) {
+            Log.d(TAG, "toCorner: startPos=" + startPos);
+            Log.d(TAG, "toCorner: finalPos=" + finalPos);
+        }
 
         ValueAnimator toCorner = ValueAnimator.ofFloat(0, 1);
         toCorner.setDuration(SCREENSHOT_TO_CORNER_Y_DURATION_MS);
@@ -427,7 +432,7 @@ public class ScreenshotView extends FrameLayout implements
             @Override
             public void onAnimationEnd(Animator animation) {
                 if (DEBUG_ANIM) {
-                    Log.d(TAG, "drop-in animation completed");
+                    Log.d(TAG, "drop-in animation ended");
                 }
                 mDismissButton.setOnClickListener(view -> {
                     if (DEBUG_INPUT) {
@@ -653,13 +658,12 @@ public class ScreenshotView extends FrameLayout implements
         getViewTreeObserver().removeOnComputeInternalInsetsListener(this);
         // Clear any references to the bitmap
         mScreenshotPreview.setImageDrawable(null);
+        mScreenshotPreview.setVisibility(View.INVISIBLE);
         mPendingSharedTransition = false;
         mActionsContainerBackground.setVisibility(View.GONE);
         mActionsContainer.setVisibility(View.GONE);
         mBackgroundProtection.setAlpha(0f);
         mDismissButton.setVisibility(View.GONE);
-        mScreenshotPreview.setVisibility(View.GONE);
-        mScreenshotPreview.setLayerType(View.LAYER_TYPE_NONE, null);
         mScreenshotStatic.setTranslationX(0);
         mScreenshotPreview.setTranslationY(0);
         mScreenshotPreview.setContentDescription(
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java
index 54b1b2c8c54c..825c85769e03 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java
@@ -72,7 +72,7 @@ public class ScrollCaptureController {
      *
      * @param after action to take after the flow is complete
      */
-    public void run(final Runnable after) {
+    public void start(final Runnable after) {
         mCaptureTime = ZonedDateTime.now();
         mRequestId = UUID.randomUUID();
         mConnection.start((session) -> startCapture(session, after));
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
index 144ad39a32aa..daa9d099de86 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
@@ -25,6 +25,7 @@ import static com.android.systemui.screenshot.LogConfig.DEBUG_DISMISS;
 import static com.android.systemui.screenshot.LogConfig.DEBUG_SERVICE;
 import static com.android.systemui.screenshot.LogConfig.logTag;
 
+import android.annotation.MainThread;
 import android.app.Service;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
@@ -59,7 +60,8 @@ import javax.inject.Inject;
 public class TakeScreenshotService extends Service {
     private static final String TAG = logTag(TakeScreenshotService.class);
 
-    private final ScreenshotController mScreenshot;
+    private ScreenshotController mScreenshot;
+
     private final UserManager mUserManager;
     private final UiEventLogger mUiEventLogger;
     private final ScreenshotNotificationsController mNotificationsController;
@@ -79,6 +81,15 @@ public class TakeScreenshotService extends Service {
         }
     };
 
+    /** Informs about coarse grained state of the Controller. */
+    interface RequestCallback {
+        /** Respond to the current request indicating the screenshot request failed.*/
+        void reportError();
+
+        /** The controller has completed handling this request UI has been removed */
+        void onFinish();
+    }
+
     @Inject
     public TakeScreenshotService(ScreenshotController screenshotController, UserManager userManager,
             UiEventLogger uiEventLogger,
@@ -116,7 +127,8 @@ public class TakeScreenshotService extends Service {
             Log.d(TAG, "onUnbind");
         }
         if (mScreenshot != null) {
-            mScreenshot.dismissScreenshot(true);
+            mScreenshot.removeWindow();
+            mScreenshot = null;
         }
         unregisterReceiver(mCloseSystemDialogs);
         return false;
@@ -126,18 +138,39 @@ public class TakeScreenshotService extends Service {
     public void onDestroy() {
         super.onDestroy();
         if (mScreenshot != null) {
+            mScreenshot.removeWindow();
             mScreenshot.releaseContext();
+            mScreenshot = null;
         }
         if (DEBUG_SERVICE) {
             Log.d(TAG, "onDestroy");
         }
     }
 
+    static class RequestCallbackImpl implements RequestCallback {
+        private final Messenger mReplyTo;
+
+        RequestCallbackImpl(Messenger replyTo) {
+            mReplyTo = replyTo;
+        }
+
+        public void reportError() {
+            reportUri(mReplyTo, null);
+            sendComplete(mReplyTo);
+        }
+
+        @Override
+        public void onFinish() {
+            sendComplete(mReplyTo);
+        }
+    }
+
     /** Respond to incoming Message via Binder (Messenger) */
+    @MainThread
     private boolean handleMessage(Message msg) {
         final Messenger replyTo = msg.replyTo;
-        final Runnable onComplete = () -> sendComplete(replyTo);
         final Consumer<Uri> uriConsumer = (uri) -> reportUri(replyTo, uri);
+        RequestCallback requestCallback = new RequestCallbackImpl(replyTo);
 
         // If the storage for this user is locked, we have no place to store
         // the screenshot, so skip taking it instead of showing a misleading
@@ -146,14 +179,7 @@ public class TakeScreenshotService extends Service {
             Log.w(TAG, "Skipping screenshot because storage is locked!");
             mNotificationsController.notifyScreenshotError(
                     R.string.screenshot_failed_to_save_user_locked_text);
-            if (DEBUG_CALLBACK) {
-                Log.d(TAG, "handleMessage: calling uriConsumer.accept(null)");
-            }
-            uriConsumer.accept(null);
-            if (DEBUG_CALLBACK) {
-                Log.d(TAG, "handleMessage: calling onComplete.run()");
-            }
-            onComplete.run();
+            requestCallback.reportError();
             return true;
         }
 
@@ -167,13 +193,13 @@ public class TakeScreenshotService extends Service {
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_FULLSCREEN");
                 }
-                mScreenshot.takeScreenshotFullscreen(uriConsumer, onComplete);
+                mScreenshot.takeScreenshotFullscreen(uriConsumer, requestCallback);
                 break;
             case WindowManager.TAKE_SCREENSHOT_SELECTED_REGION:
                 if (DEBUG_SERVICE) {
                     Log.d(TAG, "handleMessage: TAKE_SCREENSHOT_SELECTED_REGION");
                 }
-                mScreenshot.takeScreenshotPartial(uriConsumer, onComplete);
+                mScreenshot.takeScreenshotPartial(uriConsumer, requestCallback);
                 break;
             case WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE:
                 if (DEBUG_SERVICE) {
@@ -186,8 +212,16 @@ public class TakeScreenshotService extends Service {
                 int taskId = screenshotRequest.getTaskId();
                 int userId = screenshotRequest.getUserId();
                 ComponentName topComponent = screenshotRequest.getTopComponent();
-                mScreenshot.handleImageAsScreenshot(screenshot, screenBounds, insets,
-                        taskId, userId, topComponent, uriConsumer, onComplete);
+
+                if (screenshot == null) {
+                    Log.e(TAG, "Got null bitmap from screenshot message");
+                    mNotificationsController.notifyScreenshotError(
+                            R.string.screenshot_failed_to_capture_text);
+                    requestCallback.reportError();
+                } else {
+                    mScreenshot.handleImageAsScreenshot(screenshot, screenBounds, insets,
+                            taskId, userId, topComponent, uriConsumer, requestCallback);
+                }
                 break;
             default:
                 Log.w(TAG, "Invalid screenshot option: " + msg.what);
@@ -196,7 +230,7 @@ public class TakeScreenshotService extends Service {
         return true;
     };
 
-    private void sendComplete(Messenger target) {
+    private static void sendComplete(Messenger target) {
         try {
             if (DEBUG_CALLBACK) {
                 Log.d(TAG, "sendComplete: " + target);
@@ -207,7 +241,7 @@ public class TakeScreenshotService extends Service {
         }
     }
 
-    private void reportUri(Messenger target, Uri uri) {
+    private static void reportUri(Messenger target, Uri uri) {
         try {
             if (DEBUG_CALLBACK) {
                 Log.d(TAG, "reportUri: " + target + " -> " + uri);