Merge "Starting window performance tuning for binder block." into sc-dev

12 files changed, 296 insertions, 144 deletions

diff --git a/core/java/android/window/StartingWindowInfo.java b/core/java/android/window/StartingWindowInfo.java
index 08bb1a957e05..385d6cffa138 100644
--- a/core/java/android/window/StartingWindowInfo.java
+++ b/core/java/android/window/StartingWindowInfo.java
@@ -142,6 +142,12 @@ public final class StartingWindowInfo implements Parcelable {
      */
     public boolean isKeyguardOccluded = false;
 
+    /**
+     * TaskSnapshot.
+     * @hide
+     */
+    public TaskSnapshot mTaskSnapshot;
+
     public StartingWindowInfo() {
 
     }
@@ -164,6 +170,7 @@ public final class StartingWindowInfo implements Parcelable {
         dest.writeTypedObject(mainWindowLayoutParams, flags);
         dest.writeInt(splashScreenThemeResId);
         dest.writeBoolean(isKeyguardOccluded);
+        dest.writeTypedObject(mTaskSnapshot, flags);
     }
 
     void readFromParcel(@NonNull Parcel source) {
@@ -175,6 +182,7 @@ public final class StartingWindowInfo implements Parcelable {
         mainWindowLayoutParams = source.readTypedObject(WindowManager.LayoutParams.CREATOR);
         splashScreenThemeResId = source.readInt();
         isKeyguardOccluded = source.readBoolean();
+        mTaskSnapshot = source.readTypedObject(TaskSnapshot.CREATOR);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
index 516dfd0758ee..147f2e2ec846 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.startingsurface;
 
+import static android.os.Process.THREAD_PRIORITY_TOP_APP_BOOST;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.os.UserHandle.getUserHandleForUid;
 
@@ -35,6 +36,8 @@ import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.LayerDrawable;
 import android.os.Build;
+import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.Trace;
 import android.util.Slog;
 import android.view.SurfaceControl;
@@ -48,10 +51,11 @@ import com.android.launcher3.icons.IconProvider;
 import com.android.wm.shell.common.TransactionPool;
 
 import java.util.List;
+import java.util.function.Consumer;
 
 /**
  * Util class to create the view for a splash screen content.
- *
+ * Everything execute in this class should be post to mSplashscreenWorkerHandler.
  * @hide
  */
 public class SplashscreenContentDrawer {
@@ -78,6 +82,7 @@ public class SplashscreenContentDrawer {
     private int mIconEarlyExitDistance;
     private final TransactionPool mTransactionPool;
     private final SplashScreenWindowAttrs mTmpAttrs = new SplashScreenWindowAttrs();
+    private final Handler mSplashscreenWorkerHandler;
 
     SplashscreenContentDrawer(Context context, int maxAnimatableIconDuration,
             int iconExitAnimDuration, int appRevealAnimDuration, TransactionPool pool) {
@@ -87,6 +92,45 @@ public class SplashscreenContentDrawer {
         mAppRevealDuration = appRevealAnimDuration;
         mIconExitDuration = iconExitAnimDuration;
         mTransactionPool = pool;
+
+        // Initialize Splashscreen worker thread
+        // TODO(b/185288910) move it into WMShellConcurrencyModule and provide an executor to make
+        //  it easier to test stuff that happens on that thread later.
+        final HandlerThread shellSplashscreenWorkerThread =
+                new HandlerThread("wmshell.splashworker", THREAD_PRIORITY_TOP_APP_BOOST);
+        shellSplashscreenWorkerThread.start();
+        mSplashscreenWorkerHandler = shellSplashscreenWorkerThread.getThreadHandler();
+    }
+
+    /**
+     * Create a SplashScreenView object.
+     *
+     * In order to speed up the splash screen view to show on first frame, preparing the
+     * view on background thread so the view and the drawable can be create and pre-draw in
+     * parallel.
+     *
+     * @param consumer Receiving the SplashScreenView object, which will also be executed
+     *                 on splash screen thread. Note that the view can be null if failed.
+     */
+    void createContentView(Context context, int splashScreenResId, ActivityInfo info,
+            int taskId, Consumer<SplashScreenView> consumer) {
+        mSplashscreenWorkerHandler.post(() -> {
+            SplashScreenView contentView;
+            try {
+                contentView = SplashscreenContentDrawer.makeSplashscreenContent(
+                        context, splashScreenResId);
+                if (contentView == null) {
+                    Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "makeSplashScreenContentView");
+                    contentView = makeSplashScreenContentView(context, info);
+                    Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+                }
+            } catch (RuntimeException e) {
+                Slog.w(TAG, "failed creating starting window content at taskId: "
+                        + taskId, e);
+                contentView = null;
+            }
+            consumer.accept(contentView);
+        });
     }
 
     private void updateDensity() {
@@ -146,7 +190,7 @@ public class SplashscreenContentDrawer {
         }
     }
 
-    SplashScreenView makeSplashScreenContentView(Context context, ActivityInfo ai) {
+    private SplashScreenView makeSplashScreenContentView(Context context, ActivityInfo ai) {
         updateDensity();
 
         getWindowAttrs(context, mTmpAttrs);
@@ -199,7 +243,7 @@ public class SplashscreenContentDrawer {
         }
     }
 
-    static class SplashScreenWindowAttrs {
+    private static class SplashScreenWindowAttrs {
         private int mWindowBgResId = 0;
         private int mWindowBgColor = Color.TRANSPARENT;
         private Drawable mReplaceIcon = null;
@@ -271,9 +315,7 @@ public class SplashscreenContentDrawer {
                 if (DEBUG) {
                     Slog.d(TAG, "The icon is not an AdaptiveIconDrawable");
                 }
-                mFinalIconDrawable = SplashscreenIconDrawableFactory.makeIconDrawable(
-                        mIconBackground != Color.TRANSPARENT
-                        ? mIconBackground : mThemeColor, mIconDrawable, mIconSize);
+                createIconDrawable(mIconDrawable, mIconSize);
             }
             final int iconSize = mFinalIconDrawable != null ? (int) (mIconSize * mScale) : 0;
             mCachedResult = fillViewWithIcon(mContext, iconSize, mFinalIconDrawable);
@@ -283,8 +325,8 @@ public class SplashscreenContentDrawer {
 
         private void createIconDrawable(Drawable iconDrawable, int iconSize) {
             mFinalIconDrawable = SplashscreenIconDrawableFactory.makeIconDrawable(
-                    mIconBackground != Color.TRANSPARENT
-                    ? mIconBackground : mThemeColor, iconDrawable, iconSize);
+                    mIconBackground != Color.TRANSPARENT ? mIconBackground : mThemeColor,
+                    iconDrawable, iconSize, mSplashscreenWorkerHandler);
         }
 
         private boolean processAdaptiveIcon() {
@@ -399,7 +441,7 @@ public class SplashscreenContentDrawer {
         return root < 0.1;
     }
 
-    static SplashScreenView makeSplashscreenContent(Context ctx,
+    private static SplashScreenView makeSplashscreenContent(Context ctx,
             int splashscreenContentResId) {
         // doesn't support windowSplashscreenContent after S
         // TODO add an allowlist to skip some packages if needed
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenIconDrawableFactory.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenIconDrawableFactory.java
index 4196d68f5800..85845d0d9c89 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenIconDrawableFactory.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenIconDrawableFactory.java
@@ -37,6 +37,7 @@ import android.graphics.drawable.AdaptiveIconDrawable;
 import android.graphics.drawable.Animatable;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
+import android.os.Handler;
 import android.os.Trace;
 import android.util.PathParser;
 import android.window.SplashScreenView;
@@ -50,42 +51,61 @@ import com.android.internal.R;
 public class SplashscreenIconDrawableFactory {
 
     static Drawable makeIconDrawable(@ColorInt int backgroundColor,
-            @NonNull Drawable foregroundDrawable, int iconSize) {
+            @NonNull Drawable foregroundDrawable, int iconSize,
+            Handler splashscreenWorkerHandler) {
         if (foregroundDrawable instanceof Animatable) {
             return new AnimatableIconDrawable(backgroundColor, foregroundDrawable);
         } else if (foregroundDrawable instanceof AdaptiveIconDrawable) {
-            return new ImmobileIconDrawable((AdaptiveIconDrawable) foregroundDrawable, iconSize);
+            return new ImmobileIconDrawable((AdaptiveIconDrawable) foregroundDrawable, iconSize,
+                    splashscreenWorkerHandler);
         } else {
             return new ImmobileIconDrawable(new AdaptiveIconDrawable(
-                    new ColorDrawable(backgroundColor), foregroundDrawable), iconSize);
+                    new ColorDrawable(backgroundColor), foregroundDrawable), iconSize,
+                    splashscreenWorkerHandler);
         }
     }
 
     private static class ImmobileIconDrawable extends Drawable {
-        private Shader mLayersShader;
+        private boolean mCacheComplete;
         private final Paint mPaint = new Paint(Paint.ANTI_ALIAS_FLAG | Paint.DITHER_FLAG
                 | Paint.FILTER_BITMAP_FLAG);
 
-        ImmobileIconDrawable(AdaptiveIconDrawable drawable, int iconSize) {
-            cachePaint(drawable, iconSize, iconSize);
+        ImmobileIconDrawable(AdaptiveIconDrawable drawable, int iconSize,
+                Handler splashscreenWorkerHandler) {
+            splashscreenWorkerHandler.post(() -> cachePaint(drawable, iconSize, iconSize));
         }
 
         private void cachePaint(AdaptiveIconDrawable drawable, int width, int height) {
-            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "cachePaint");
-            final Bitmap layersBitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
-            final Canvas canvas = new Canvas(layersBitmap);
-            drawable.setBounds(0, 0, width, height);
-            drawable.draw(canvas);
-            mLayersShader = new BitmapShader(layersBitmap, Shader.TileMode.CLAMP,
-                    Shader.TileMode.CLAMP);
-            mPaint.setShader(mLayersShader);
-            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+            synchronized (mPaint) {
+                if (mCacheComplete) {
+                    return;
+                }
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "cachePaint");
+                final Bitmap layersBitmap = Bitmap.createBitmap(width, height,
+                        Bitmap.Config.ARGB_8888);
+                final Canvas canvas = new Canvas(layersBitmap);
+                drawable.setBounds(0, 0, width, height);
+                drawable.draw(canvas);
+                final Shader layersShader = new BitmapShader(layersBitmap, Shader.TileMode.CLAMP,
+                        Shader.TileMode.CLAMP);
+                mPaint.setShader(layersShader);
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+                mCacheComplete = true;
+            }
         }
 
         @Override
         public void draw(Canvas canvas) {
-            final Rect bounds = getBounds();
-            canvas.drawRect(bounds, mPaint);
+            synchronized (mPaint) {
+                if (mCacheComplete) {
+                    final Rect bounds = getBounds();
+                    canvas.drawRect(bounds, mPaint);
+                } else {
+                    // this shouldn't happen, but if it really happen, invalidate self to wait
+                    // for cachePaint finish.
+                    invalidateSelf();
+                }
+            }
         }
 
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
index e95135a1a59f..7037d18decbe 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawer.java
@@ -18,8 +18,10 @@ package com.android.wm.shell.startingsurface;
 
 import static android.content.Context.CONTEXT_RESTRICTED;
 import static android.content.res.Configuration.EMPTY;
+import static android.view.Choreographer.CALLBACK_INSETS_ANIMATION;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import android.annotation.Nullable;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.ActivityTaskManager;
 import android.content.Context;
@@ -29,16 +31,15 @@ import android.content.pm.PackageManager;
 import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
-import android.graphics.Color;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
-import android.graphics.drawable.ColorDrawable;
 import android.hardware.display.DisplayManager;
 import android.os.IBinder;
 import android.os.SystemClock;
 import android.os.UserHandle;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.view.Choreographer;
 import android.view.Display;
 import android.view.SurfaceControl;
 import android.view.View;
@@ -54,9 +55,44 @@ import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
 
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 
 /**
  * A class which able to draw splash screen or snapshot as the starting window for a task.
+ *
+ * In order to speed up, there will use two threads to creating a splash screen in parallel.
+ * Right now we are still using PhoneWindow to create splash screen window, so the view is added to
+ * the ViewRootImpl, and those view won't be draw immediately because the ViewRootImpl will call
+ * scheduleTraversal to register a callback from Choreographer, so the drawing result of the view
+ * can synchronize on each frame.
+ *
+ * The bad thing is that we cannot decide when would Choreographer#doFrame happen, and drawing
+ * the AdaptiveIconDrawable object can be time consuming, so we use the splash-screen background
+ * thread to draw the AdaptiveIconDrawable object to a Bitmap and cache it to a BitmapShader after
+ * the SplashScreenView just created, once we get the BitmapShader then the #draw call can be very
+ * quickly.
+ *
+ * So basically we are using the spare time to prepare the SplashScreenView while splash screen
+ * thread is waiting for
+ * 1. WindowManager#addView(binder call to WM),
+ * 2. Choreographer#doFrame happen(uncertain time for next frame, depends on device),
+ * 3. Session#relayout(another binder call to WM which under Choreographer#doFrame, but will
+ * always happen before #draw).
+ * Because above steps are running on splash-screen thread, so pre-draw the BitmapShader on
+ * splash-screen background tread can make they execute in parallel, which ensure it is faster then
+ * to draw the AdaptiveIconDrawable when receive callback from Choreographer#doFrame.
+ *
+ * Here is the sequence to compare the difference between using single and two thread.
+ *
+ * Single thread:
+ * => makeSplashScreenContentView -> WM#addView .. waiting for Choreographer#doFrame -> relayout
+ * -> draw -> AdaptiveIconDrawable#draw
+ *
+ * Two threads:
+ * => makeSplashScreenContentView -> cachePaint(=AdaptiveIconDrawable#draw)
+ * => WM#addView -> .. waiting for Choreographer#doFrame -> relayout -> draw -> (draw the Paint
+ * directly).
+ *
  * @hide
  */
 public class StartingSurfaceDrawer {
@@ -68,7 +104,11 @@ public class StartingSurfaceDrawer {
     private final DisplayManager mDisplayManager;
     private final ShellExecutor mSplashScreenExecutor;
     private final SplashscreenContentDrawer mSplashscreenContentDrawer;
+    private Choreographer mChoreographer;
 
+    /**
+     * @param splashScreenExecutor The thread used to control add and remove starting window.
+     */
     public StartingSurfaceDrawer(Context context, ShellExecutor splashScreenExecutor,
             TransactionPool pool) {
         mContext = context;
@@ -82,6 +122,7 @@ public class StartingSurfaceDrawer {
                 com.android.wm.shell.R.integer.starting_window_app_reveal_anim_duration);
         mSplashscreenContentDrawer = new SplashscreenContentDrawer(mContext,
                 maxAnimatableIconDuration, iconExitAnimDuration, appRevealAnimDuration, pool);
+        mSplashScreenExecutor.execute(() -> mChoreographer = Choreographer.getInstance());
     }
 
     private final SparseArray<StartingWindowRecord> mStartingWindowRecords = new SparseArray<>();
@@ -267,36 +308,91 @@ public class StartingSurfaceDrawer {
         params.setTitle("Splash Screen " + activityInfo.packageName);
 
         // TODO(b/173975965) tracking performance
-        SplashScreenView sView = null;
+        // Prepare the splash screen content view on splash screen worker thread in parallel, so the
+        // content view won't be blocked by binder call like addWindow and relayout.
+        // 1. Trigger splash screen worker thread to create SplashScreenView before/while
+        // Session#addWindow.
+        // 2. Synchronize the SplashscreenView to splash screen thread before Choreographer start
+        // traversal, which will call Session#relayout on splash screen thread.
+        // 3. Pre-draw the BitmapShader if the icon is immobile on splash screen worker thread, at
+        // the same time the splash screen thread should be executing Session#relayout. Blocking the
+        // traversal -> draw on splash screen thread until the BitmapShader of the icon is ready.
+        final Runnable setViewSynchronized;
+        if (!emptyView) {
+            // Record whether create splash screen view success, notify to current thread after
+            // create splash screen view finished.
+            final SplashScreenViewSupplier viewSupplier = new SplashScreenViewSupplier();
+            setViewSynchronized = () -> {
+                // waiting for setContentView before relayoutWindow
+                SplashScreenView contentView = viewSupplier.get();
+                final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
+                // if record == null, either the starting window added fail or removed already.
+                if (record != null) {
+                    // if view == null then creation of content view was failed.
+                    if (contentView != null) {
+                        try {
+                            win.setContentView(contentView);
+                            contentView.cacheRootWindow(win);
+                        } catch (RuntimeException e) {
+                            Slog.w(TAG, "failed set content view to starting window "
+                                    + "at taskId: " + taskId, e);
+                            contentView = null;
+                        }
+                    }
+                    record.setSplashScreenView(contentView);
+                }
+            };
+            mSplashscreenContentDrawer.createContentView(context,
+                    splashscreenContentResId[0], activityInfo, taskId, viewSupplier::setView);
+        } else {
+            setViewSynchronized = null;
+        }
+
         try {
             final View view = win.getDecorView();
             final WindowManager wm = mContext.getSystemService(WindowManager.class);
-            if (!emptyView) {
-                // splash screen content will be deprecated after S.
-                sView = SplashscreenContentDrawer.makeSplashscreenContent(
-                        context, splashscreenContentResId[0]);
-                final boolean splashscreenContentCompatible = sView != null;
-                if (splashscreenContentCompatible) {
-                    win.setContentView(sView);
-                } else {
-                    sView = mSplashscreenContentDrawer
-                            .makeSplashScreenContentView(context, activityInfo);
-                    win.setContentView(sView);
-                    win.setBackgroundDrawable(new ColorDrawable(Color.TRANSPARENT));
-                    sView.cacheRootWindow(win);
-                }
-            }
             postAddWindow(taskId, appToken, view, wm, params);
+
+            // all done
+            if (emptyView) {
+                return;
+            }
+            // We use the splash screen worker thread to create SplashScreenView while adding the
+            // window, as otherwise Choreographer#doFrame might be delayed on this thread.
+            // And since Choreographer#doFrame won't happen immediately after adding the window, if
+            // the view is not added to the PhoneWindow on the first #doFrame, the view will not be
+            // rendered on the first frame. So here we need to synchronize the view on the window
+            // before first round relayoutWindow, which will happen after insets animation.
+            mChoreographer.postCallback(CALLBACK_INSETS_ANIMATION, setViewSynchronized, null);
         } catch (RuntimeException e) {
             // don't crash if something else bad happens, for example a
             // failure loading resources because we are loading from an app
             // on external storage that has been unmounted.
-            Slog.w(TAG, " failed creating starting window at taskId: " + taskId, e);
-            sView = null;
-        } finally {
-            final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
-            if (record != null) {
-                record.setSplashScreenView(sView);
+            Slog.w(TAG, "failed creating starting window at taskId: " + taskId, e);
+        }
+    }
+
+    private static class SplashScreenViewSupplier implements Supplier<SplashScreenView> {
+        private SplashScreenView mView;
+        private boolean mIsViewSet;
+        void setView(SplashScreenView view) {
+            synchronized (this) {
+                mView = view;
+                mIsViewSet = true;
+                notify();
+            }
+        }
+
+        @Override
+        public @Nullable SplashScreenView get() {
+            synchronized (this) {
+                while (!mIsViewSet) {
+                    try {
+                        wait();
+                    } catch (InterruptedException ignored) {
+                    }
+                }
+                return mView;
             }
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
index 8a629bcf4f3b..e3362870cdf0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
@@ -16,6 +16,7 @@
 package com.android.wm.shell.startingsurface;
 
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.window.StartingWindowInfo.STARTING_WINDOW_TYPE_EMPTY_SPLASH_SCREEN;
 import static android.window.StartingWindowInfo.STARTING_WINDOW_TYPE_NONE;
 import static android.window.StartingWindowInfo.STARTING_WINDOW_TYPE_SNAPSHOT;
@@ -30,11 +31,11 @@ import static android.window.StartingWindowInfo.TYPE_PARAMETER_TASK_SWITCH;
 import static com.android.wm.shell.common.ExecutorUtils.executeRemoteCallWithTaskPermission;
 
 import android.app.ActivityManager.RunningTaskInfo;
-import android.app.ActivityTaskManager;
 import android.content.Context;
 import android.graphics.Rect;
 import android.os.IBinder;
 import android.os.RemoteException;
+import android.os.Trace;
 import android.util.Slog;
 import android.view.SurfaceControl;
 import android.window.StartingWindowInfo;
@@ -109,17 +110,9 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
     }
 
     private static class StartingTypeChecker {
-        TaskSnapshot mSnapshot;
-
-        StartingTypeChecker() { }
-
-        private void reset() {
-            mSnapshot = null;
-        }
 
         private @StartingWindowInfo.StartingWindowType int
                 estimateStartingWindowType(StartingWindowInfo windowInfo) {
-            reset();
             final int parameter = windowInfo.startingWindowTypeParameter;
             final boolean newTask = (parameter & TYPE_PARAMETER_NEW_TASK) != 0;
             final boolean taskSwitch = (parameter & TYPE_PARAMETER_TASK_SWITCH) != 0;
@@ -159,7 +152,7 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
                 }
             }
             if (taskSwitch && allowTaskSnapshot) {
-                final TaskSnapshot snapshot = getTaskSnapshot(windowInfo.taskInfo.taskId);
+                final TaskSnapshot snapshot = windowInfo.mTaskSnapshot;
                 if (isSnapshotCompatible(windowInfo, snapshot)) {
                     return STARTING_WINDOW_TYPE_SNAPSHOT;
                 }
@@ -198,20 +191,6 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
             }
             return taskRotation == snapshotRotation;
         }
-
-        private TaskSnapshot getTaskSnapshot(int taskId) {
-            if (mSnapshot != null) {
-                return mSnapshot;
-            }
-            try {
-                mSnapshot = ActivityTaskManager.getService().getTaskSnapshot(taskId,
-                        false/* isLowResolution */);
-            } catch (RemoteException e) {
-                Slog.e(TAG, "Unable to get snapshot for task: " + taskId + ", from: " + e);
-                return null;
-            }
-            return mSnapshot;
-        }
     }
 
     /*
@@ -232,7 +211,9 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
      */
     public void addStartingWindow(StartingWindowInfo windowInfo, IBinder appToken) {
         mSplashScreenExecutor.execute(() -> {
-            final int suggestionType = mStartingTypeChecker.estimateStartingWindowType(windowInfo);
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "addStartingWindow");
+            final int suggestionType = mStartingTypeChecker.estimateStartingWindowType(
+                    windowInfo);
             final RunningTaskInfo runningTaskInfo = windowInfo.taskInfo;
             if (mTaskLaunchingCallback != null && shouldSendToListener(suggestionType)) {
                 mTaskLaunchingCallback.accept(runningTaskInfo.taskId, suggestionType);
@@ -244,10 +225,12 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
                 mStartingSurfaceDrawer.addSplashScreenStartingWindow(windowInfo, appToken,
                         true /* emptyView */);
             } else if (suggestionType == STARTING_WINDOW_TYPE_SNAPSHOT) {
-                final TaskSnapshot snapshot = mStartingTypeChecker.mSnapshot;
-                mStartingSurfaceDrawer.makeTaskSnapshotWindow(windowInfo, appToken, snapshot);
+                final TaskSnapshot snapshot = windowInfo.mTaskSnapshot;
+                mStartingSurfaceDrawer.makeTaskSnapshotWindow(windowInfo, appToken,
+                        snapshot);
             }
             // If prefer don't show, then don't show!
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         });
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
index 6e437411dcf8..e9ce2addd464 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/TaskSnapshotWindow.java
@@ -19,6 +19,7 @@ package com.android.wm.shell.startingsurface;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
 import static android.graphics.Color.WHITE;
 import static android.graphics.Color.alpha;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.view.WindowInsetsController.APPEARANCE_LIGHT_NAVIGATION_BARS;
 import static android.view.WindowInsetsController.APPEARANCE_LIGHT_STATUS_BARS;
 import static android.view.WindowManager.LayoutParams.FLAG_ALT_FOCUSABLE_IM;
@@ -60,6 +61,7 @@ import android.hardware.HardwareBuffer;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.os.Trace;
 import android.util.MergedConfiguration;
 import android.util.Slog;
 import android.view.IWindowSession;
@@ -222,8 +224,10 @@ public class TaskSnapshotWindow {
         final InputChannel tmpInputChannel = new InputChannel();
         mainExecutor.execute(() -> {
             try {
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "TaskSnapshot#addToDisplay");
                 final int res = session.addToDisplay(window, layoutParams, View.GONE, displayId,
                         mTmpInsetsState, tmpInputChannel, mTmpInsetsState, mTempControls);
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
                 if (res < 0) {
                     Slog.w(TAG, "Failed to add snapshot starting window res=" + res);
                     return;
@@ -233,9 +237,11 @@ public class TaskSnapshotWindow {
             }
             window.setOuter(snapshotSurface);
             try {
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "TaskSnapshot#relayout");
                 session.relayout(window, layoutParams, -1, -1, View.VISIBLE, 0, -1,
                         tmpFrames, tmpMergedConfiguration, surfaceControl, mTmpInsetsState,
                         mTempControls, TMP_SURFACE_SIZE);
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
             } catch (RemoteException e) {
                 snapshotSurface.clearWindowSynced();
             }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
index d1d13131dd44..262353502bb7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingSurfaceDrawerTests.java
@@ -77,9 +77,9 @@ public class StartingSurfaceDrawerTests {
         int mAddWindowForTask = 0;
         int mViewThemeResId;
 
-        TestStartingSurfaceDrawer(Context context, ShellExecutor animExecutor,
+        TestStartingSurfaceDrawer(Context context, ShellExecutor splashScreenExecutor,
                 TransactionPool pool) {
-            super(context, animExecutor, pool);
+            super(context, splashScreenExecutor, pool);
         }
 
         @Override
@@ -120,8 +120,9 @@ public class StartingSurfaceDrawerTests {
         doReturn(metrics).when(mMockWindowManager).getMaximumWindowMetrics();
         doNothing().when(mMockWindowManager).addView(any(), any());
 
-        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(context,
-                new HandlerExecutor(new Handler(Looper.getMainLooper())),
+        final HandlerExecutor testExecutor =
+                new HandlerExecutor(new Handler(Looper.getMainLooper()));
+        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(context, testExecutor,
                 mTransactionPool));
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
index 3fc3d89461ae..f96d34495fda 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShellBaseModule.java
@@ -381,8 +381,8 @@ public abstract class WMShellBaseModule {
     @WMSingleton
     @Provides
     static StartingWindowController provideStartingWindowController(Context context,
-            @ShellSplashscreenThread ShellExecutor executor, TransactionPool pool) {
-        return new StartingWindowController(context, executor, pool);
+            @ShellSplashscreenThread ShellExecutor splashScreenExecutor, TransactionPool pool) {
+        return new StartingWindowController(context, splashScreenExecutor, pool);
     }
 
     //
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index aa9727a28985..c39358e7ad4a 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -1926,12 +1926,16 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
     }
 
     void scheduleAddStartingWindow() {
-        // Note: we really want to do sendMessageAtFrontOfQueue() because we
-        // want to process the message ASAP, before any other queued
-        // messages.
-        if (!mWmService.mAnimationHandler.hasCallbacks(mAddStartingWindow)) {
-            ProtoLog.v(WM_DEBUG_STARTING_WINDOW, "Enqueueing ADD_STARTING");
-            mWmService.mAnimationHandler.postAtFrontOfQueue(mAddStartingWindow);
+        if (StartingSurfaceController.DEBUG_ENABLE_SHELL_DRAWER) {
+            mAddStartingWindow.run();
+        } else {
+            // Note: we really want to do sendMessageAtFrontOfQueue() because we
+            // want to process the message ASAP, before any other queued
+            // messages.
+            if (!mWmService.mAnimationHandler.hasCallbacks(mAddStartingWindow)) {
+                ProtoLog.v(WM_DEBUG_STARTING_WINDOW, "Enqueueing ADD_STARTING");
+                mWmService.mAnimationHandler.postAtFrontOfQueue(mAddStartingWindow);
+            }
         }
     }
 
@@ -1943,7 +1947,9 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
             final StartingData startingData;
             synchronized (mWmService.mGlobalLock) {
                 // There can only be one adding request, silly caller!
-                mWmService.mAnimationHandler.removeCallbacks(this);
+                if (!StartingSurfaceController.DEBUG_ENABLE_SHELL_DRAWER) {
+                    mWmService.mAnimationHandler.removeCallbacks(this);
+                }
 
                 if (mStartingData == null) {
                     // Animation has been canceled... do nothing.
@@ -2192,17 +2198,22 @@ final class ActivityRecord extends WindowToken implements WindowManagerService.A
                 + " startingView=%s Callers=%s", this, mStartingWindow, mStartingSurface,
                 Debug.getCallers(5));
 
-
-        // Use the same thread to remove the window as we used to add it, as otherwise we end up
-        // with things in the view hierarchy being called from different threads.
-        mWmService.mAnimationHandler.post(() -> {
+        final Runnable removeSurface = () -> {
             ProtoLog.v(WM_DEBUG_STARTING_WINDOW, "Removing startingView=%s", surface);
             try {
                 surface.remove(prepareAnimation);
             } catch (Exception e) {
                 Slog.w(TAG_WM, "Exception when removing starting window", e);
             }
-        });
+        };
+
+        if (StartingSurfaceController.DEBUG_ENABLE_SHELL_DRAWER) {
+            removeSurface.run();
+        } else {
+            // Use the same thread to remove the window as we used to add it, as otherwise we end up
+            // with things in the view hierarchy being called from different threads.
+            mWmService.mAnimationHandler.post(removeSurface);
+        }
     }
 
     private void removeAppTokenFromDisplay() {
diff --git a/services/core/java/com/android/server/wm/StartingSurfaceController.java b/services/core/java/com/android/server/wm/StartingSurfaceController.java
index 603bfd153481..a9b06ca5042b 100644
--- a/services/core/java/com/android/server/wm/StartingSurfaceController.java
+++ b/services/core/java/com/android/server/wm/StartingSurfaceController.java
@@ -61,7 +61,7 @@ public class StartingSurfaceController {
         synchronized (mService.mGlobalLock) {
             final Task task = activity.getTask();
             if (task != null && mService.mAtmService.mTaskOrganizerController.addStartingWindow(
-                    task, activity.token, theme)) {
+                    task, activity.token, theme, null /* taskSnapshot */)) {
                 return new ShellStartingSurface(task);
             }
         }
@@ -128,7 +128,7 @@ public class StartingSurfaceController {
             }
             if (DEBUG_ENABLE_SHELL_DRAWER) {
                 mService.mAtmService.mTaskOrganizerController.addStartingWindow(task,
-                        activity.token, 0 /* launchTheme */);
+                        activity.token, 0 /* launchTheme */, taskSnapshot);
                 return new ShellStartingSurface(task);
             }
         }
diff --git a/services/core/java/com/android/server/wm/TaskOrganizerController.java b/services/core/java/com/android/server/wm/TaskOrganizerController.java
index 565804fc9c8e..ccc0916cf25b 100644
--- a/services/core/java/com/android/server/wm/TaskOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskOrganizerController.java
@@ -43,6 +43,7 @@ import android.window.ITaskOrganizer;
 import android.window.ITaskOrganizerController;
 import android.window.StartingWindowInfo;
 import android.window.TaskAppearedInfo;
+import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -118,25 +119,25 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
             return mTaskOrganizer.asBinder();
         }
 
-        void addStartingWindow(Task task, IBinder appToken, int launchTheme) {
+        void addStartingWindow(Task task, IBinder appToken, int launchTheme,
+                TaskSnapshot taskSnapshot) {
             final StartingWindowInfo info = task.getStartingWindowInfo();
             if (launchTheme != 0) {
                 info.splashScreenThemeResId = launchTheme;
             }
-            mDeferTaskOrgCallbacksConsumer.accept(() -> {
-                try {
-                    mTaskOrganizer.addStartingWindow(info, appToken);
-                } catch (RemoteException e) {
-                    Slog.e(TAG, "Exception sending onTaskStart callback", e);
-                }
-            });
+            info.mTaskSnapshot = taskSnapshot;
+            // make this happen prior than prepare surface
+            try {
+                mTaskOrganizer.addStartingWindow(info, appToken);
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Exception sending onTaskStart callback", e);
+            }
         }
 
         void removeStartingWindow(Task task, boolean prepareAnimation) {
-            mDeferTaskOrgCallbacksConsumer.accept(() -> {
-                SurfaceControl firstWindowLeash = null;
-                Rect mainFrame = null;
-                // TODO enable shift up animation once we fix flicker test
+            SurfaceControl firstWindowLeash = null;
+            Rect mainFrame = null;
+            // TODO enable shift up animation once we fix flicker test
 //                final boolean playShiftUpAnimation = !task.inMultiWindowMode();
 //                if (prepareAnimation && playShiftUpAnimation) {
 //                    final ActivityRecord topActivity = task.topActivityWithStartingWindow();
@@ -144,32 +145,29 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
 //                        final WindowState mainWindow =
 //                                topActivity.findMainWindow(false/* includeStartingApp */);
 //                        if (mainWindow != null) {
-                // TODO create proper leash instead of the copied SC
+            // TODO create proper leash instead of the copied SC
 //                            firstWindowLeash = new SurfaceControl(mainWindow.getSurfaceControl(),
 //                                    "TaskOrganizerController.removeStartingWindow");
 //                            mainFrame = mainWindow.getRelativeFrame();
 //                        }
 //                    }
 //                }
-                try {
-                    mTaskOrganizer.removeStartingWindow(task.mTaskId, firstWindowLeash, mainFrame,
-                    /* TODO(183004107) Revert this when jankiness is solved
-                        prepareAnimation); */ false);
+            try {
+                mTaskOrganizer.removeStartingWindow(task.mTaskId, firstWindowLeash, mainFrame,
+                /* TODO(183004107) Revert this when jankiness is solved
+                    prepareAnimation); */ false);
 
-                } catch (RemoteException e) {
-                    Slog.e(TAG, "Exception sending onStartTaskFinished callback", e);
-                }
-            });
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Exception sending onStartTaskFinished callback", e);
+            }
         }
 
         void copySplashScreenView(Task task) {
-            mDeferTaskOrgCallbacksConsumer.accept(() -> {
-                try {
-                    mTaskOrganizer.copySplashScreenView(task.mTaskId);
-                } catch (RemoteException e) {
-                    Slog.e(TAG, "Exception sending copyStartingWindowView callback", e);
-                }
-            });
+            try {
+                mTaskOrganizer.copySplashScreenView(task.mTaskId);
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Exception sending copyStartingWindowView callback", e);
+            }
         }
 
         SurfaceControl prepareLeash(Task task, boolean visible, String reason) {
@@ -266,8 +264,9 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
             mUid = uid;
         }
 
-        void addStartingWindow(Task t, IBinder appToken, int launchTheme) {
-            mOrganizer.addStartingWindow(t, appToken, launchTheme);
+        void addStartingWindow(Task t, IBinder appToken, int launchTheme,
+                TaskSnapshot taskSnapshot) {
+            mOrganizer.addStartingWindow(t, appToken, launchTheme, taskSnapshot);
         }
 
         void removeStartingWindow(Task t, boolean prepareAnimation) {
@@ -505,14 +504,15 @@ class TaskOrganizerController extends ITaskOrganizerController.Stub {
         return !ArrayUtils.contains(UNSUPPORTED_WINDOWING_MODES, winMode);
     }
 
-    boolean addStartingWindow(Task task, IBinder appToken, int launchTheme) {
+    boolean addStartingWindow(Task task, IBinder appToken, int launchTheme,
+            TaskSnapshot taskSnapshot) {
         final Task rootTask = task.getRootTask();
         if (rootTask == null || rootTask.mTaskOrganizer == null) {
             return false;
         }
         final TaskOrganizerState state =
                 mTaskOrganizerStates.get(rootTask.mTaskOrganizer.asBinder());
-        state.addStartingWindow(task, appToken, launchTheme);
+        state.addStartingWindow(task, appToken, launchTheme, taskSnapshot);
         return true;
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 124f6ddc8391..8cede6df491d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -2358,21 +2358,6 @@ public class ActivityRecordTests extends WindowTestsBase {
     }
 
     @Test
-    public void testAddRemoveRace() {
-        registerTestStartingWindowOrganizer();
-        // There was once a race condition between adding and removing starting windows
-        final ActivityRecord appToken = new ActivityBuilder(mAtm).setCreateTask(true).build();
-        for (int i = 0; i < 1000; i++) {
-            appToken.addStartingWindow(mPackageName,
-                    android.R.style.Theme, null, "Test", 0, 0, 0, 0, null, true, true, false, true,
-                    false, false);
-            appToken.removeStartingWindow();
-            waitUntilHandlersIdle();
-            assertNoStartingWindow(appToken);
-        }
-    }
-
-    @Test
     public void testTransferStartingWindow() {
         registerTestStartingWindowOrganizer();
         final ActivityRecord activity1 = new ActivityBuilder(mAtm).setCreateTask(true).build();