Cache window and icon color of splash screen starting window

The background color of splash screen starting window may depend
on the declared window background and the content if icon. That
may be expensive operations if the drawables are not using pure
color. Because the attributes to compute the colors are simple,
it only need a few fields to identify whether the incoming request
has been computed. So it is worth to cache them with very small
extra memory. The cache of a package will be cleared if the package
is removed or replaced.

It may save the creation time up to 20ms (depends on the complexity
of the drawable).

Bug: 190176296
Test: atest StartingSurfaceDrawerTests#testColorCache
Test: Cold/warm launch app twice and compare the duration of
      trace name "makeSplashScreenContentView". The second time
      should be faster. Re-install the app and launch again,
      the trace will contain "peekWindowBGColor" that means
      that cache for the package was cleared.

Change-Id: I61158416694a226999843fa8fc6a8b23d6893f02

3 files changed, 272 insertions, 91 deletions

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 1d37a128da0c..51a67e213144 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
@@ -22,7 +22,10 @@ import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import android.annotation.ColorInt;
 import android.annotation.NonNull;
 import android.app.ActivityThread;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
@@ -34,15 +37,19 @@ import android.graphics.drawable.AdaptiveIconDrawable;
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.LayerDrawable;
+import android.net.Uri;
 import android.os.Handler;
 import android.os.HandlerThread;
 import android.os.Trace;
+import android.os.UserHandle;
+import android.util.ArrayMap;
 import android.util.Slog;
 import android.view.SurfaceControl;
 import android.view.View;
 import android.window.SplashScreenView;
 
 import com.android.internal.R;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.graphics.palette.Palette;
 import com.android.internal.graphics.palette.Quantizer;
 import com.android.internal.graphics.palette.VariationalKMeansQuantizer;
@@ -51,6 +58,8 @@ import com.android.wm.shell.common.TransactionPool;
 
 import java.util.List;
 import java.util.function.Consumer;
+import java.util.function.IntSupplier;
+import java.util.function.Supplier;
 
 /**
  * Util class to create the view for a splash screen content.
@@ -76,9 +85,12 @@ public class SplashscreenContentDrawer {
     private int mBrandingImageWidth;
     private int mBrandingImageHeight;
     private int mMainWindowShiftLength;
+    private int mLastPackageContextConfigHash;
     private final TransactionPool mTransactionPool;
     private final SplashScreenWindowAttrs mTmpAttrs = new SplashScreenWindowAttrs();
     private final Handler mSplashscreenWorkerHandler;
+    @VisibleForTesting
+    final ColorCache mColorCache;
 
     SplashscreenContentDrawer(Context context, TransactionPool pool) {
         mContext = context;
@@ -92,6 +104,7 @@ public class SplashscreenContentDrawer {
                 new HandlerThread("wmshell.splashworker", THREAD_PRIORITY_TOP_APP_BOOST);
         shellSplashscreenWorkerThread.start();
         mSplashscreenWorkerHandler = shellSplashscreenWorkerThread.getThreadHandler();
+        mColorCache = new ColorCache(mContext, mSplashscreenWorkerHandler);
     }
 
     /**
@@ -183,14 +196,14 @@ public class SplashscreenContentDrawer {
         updateDensity();
 
         getWindowAttrs(context, mTmpAttrs);
-        final StartingWindowViewBuilder builder = new StartingWindowViewBuilder();
-        final int themeBGColor = peekWindowBGColor(context, this.mTmpAttrs);
+        mLastPackageContextConfigHash = context.getResources().getConfiguration().hashCode();
+        final int themeBGColor = mColorCache.getWindowColor(ai.packageName,
+                mLastPackageContextConfigHash, mTmpAttrs.mWindowBgColor, mTmpAttrs.mWindowBgResId,
+                () -> peekWindowBGColor(context, mTmpAttrs)).mBgColor;
         // TODO (b/173975965) Tracking the performance on improved splash screen.
-        return builder
-                .setContext(context)
+        return new StartingWindowViewBuilder(context, ai)
                 .setWindowBGColor(themeBGColor)
                 .makeEmptyView(emptyView)
-                .setActivityInfo(ai)
                 .build();
     }
 
@@ -232,50 +245,30 @@ public class SplashscreenContentDrawer {
     }
 
     private class StartingWindowViewBuilder {
-        private ActivityInfo mActivityInfo;
-        private Context mContext;
-        private boolean mEmptyView;
+        private final Context mContext;
+        private final ActivityInfo mActivityInfo;
 
-        // result
-        private boolean mBuildComplete = false;
-        private SplashScreenView mCachedResult;
+        private boolean mEmptyView;
         private int mThemeColor;
         private Drawable mFinalIconDrawable;
         private int mFinalIconSize = mIconSize;
 
+        StartingWindowViewBuilder(@NonNull Context context, @NonNull ActivityInfo aInfo) {
+            mContext = context;
+            mActivityInfo = aInfo;
+        }
+
         StartingWindowViewBuilder setWindowBGColor(@ColorInt int background) {
             mThemeColor = background;
-            mBuildComplete = false;
             return this;
         }
 
         StartingWindowViewBuilder makeEmptyView(boolean empty) {
             mEmptyView = empty;
-            mBuildComplete = false;
-            return this;
-        }
-
-        StartingWindowViewBuilder setActivityInfo(ActivityInfo ai) {
-            mActivityInfo = ai;
-            mBuildComplete = false;
-            return this;
-        }
-
-        StartingWindowViewBuilder setContext(Context context) {
-            mContext = context;
-            mBuildComplete = false;
             return this;
         }
 
         SplashScreenView build() {
-            if (mBuildComplete) {
-                return mCachedResult;
-            }
-            if (mContext == null || mActivityInfo == null) {
-                Slog.e(TAG, "Unable to create StartingWindowView, lack of materials!");
-                return null;
-            }
-
             Drawable iconDrawable;
             final int animationDuration;
             if (mEmptyView) {
@@ -292,7 +285,9 @@ public class SplashscreenContentDrawer {
                 final int densityDpi = mContext.getResources().getDisplayMetrics().densityDpi;
                 final int scaledIconDpi =
                         (int) (0.5f + iconScale * densityDpi * NO_BACKGROUND_SCALE);
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "getIcon");
                 iconDrawable = mIconProvider.getIcon(mActivityInfo, scaledIconDpi);
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
                 if (iconDrawable == null) {
                     iconDrawable = mContext.getPackageManager().getDefaultActivityIcon();
                 }
@@ -306,9 +301,7 @@ public class SplashscreenContentDrawer {
                 animationDuration = 0;
             }
 
-            mCachedResult = fillViewWithIcon(mFinalIconSize, mFinalIconDrawable, animationDuration);
-            mBuildComplete = true;
-            return mCachedResult;
+            return fillViewWithIcon(mFinalIconSize, mFinalIconDrawable, animationDuration);
         }
 
         private void createIconDrawable(Drawable iconDrawable, boolean legacy) {
@@ -331,28 +324,20 @@ public class SplashscreenContentDrawer {
             }
 
             Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "processAdaptiveIcon");
-            final AdaptiveIconDrawable adaptiveIconDrawable =
-                    (AdaptiveIconDrawable) iconDrawable;
-            final DrawableColorTester backIconTester =
-                    new DrawableColorTester(adaptiveIconDrawable.getBackground());
-
+            final AdaptiveIconDrawable adaptiveIconDrawable = (AdaptiveIconDrawable) iconDrawable;
             final Drawable iconForeground = adaptiveIconDrawable.getForeground();
-            final DrawableColorTester foreIconTester =
-                    new DrawableColorTester(iconForeground, true /* filterTransparent */);
-
-            final boolean foreComplex = foreIconTester.isComplexColor();
-            final int foreMainColor = foreIconTester.getDominateColor();
+            final ColorCache.IconColor iconColor = mColorCache.getIconColor(
+                    mActivityInfo.packageName, mActivityInfo.getIconResource(),
+                    mLastPackageContextConfigHash,
+                    () -> new DrawableColorTester(iconForeground, true /* filterTransparent */),
+                    () -> new DrawableColorTester(adaptiveIconDrawable.getBackground()));
 
             if (DEBUG) {
-                Slog.d(TAG, "foreground complex color? " + foreComplex + " main color: "
-                        + Integer.toHexString(foreMainColor));
-            }
-            final boolean backComplex = backIconTester.isComplexColor();
-            final int backMainColor = backIconTester.getDominateColor();
-            if (DEBUG) {
-                Slog.d(TAG, "background complex color? " + backComplex + " main color: "
-                        + Integer.toHexString(backMainColor));
-                Slog.d(TAG, "theme color? " + Integer.toHexString(mThemeColor));
+                Slog.d(TAG, "FgMainColor=" + Integer.toHexString(iconColor.mFgColor)
+                        + " BgMainColor=" + Integer.toHexString(iconColor.mBgColor)
+                        + " IsBgComplex=" + iconColor.mIsBgComplex
+                        + " FromCache=" + (iconColor.mReuseCount > 0)
+                        + " ThemeColor=" + Integer.toHexString(mThemeColor));
             }
 
             // Only draw the foreground of AdaptiveIcon to the splash screen if below condition
@@ -363,17 +348,17 @@ public class SplashscreenContentDrawer {
             // C. The background of the adaptive icon is grayscale, and the foreground of the
             // adaptive icon forms a certain contrast with the theme color.
             // D. Didn't specify icon background color.
-            if (!backComplex && mTmpAttrs.mIconBgColor == Color.TRANSPARENT
-                    && (isRgbSimilarInHsv(mThemeColor, backMainColor)
-                    || (backIconTester.isGrayscale()
-                    && !isRgbSimilarInHsv(mThemeColor, foreMainColor)))) {
+            if (!iconColor.mIsBgComplex && mTmpAttrs.mIconBgColor == Color.TRANSPARENT
+                    && (isRgbSimilarInHsv(mThemeColor, iconColor.mBgColor)
+                            || (iconColor.mIsBgGrayscale
+                                    && !isRgbSimilarInHsv(mThemeColor, iconColor.mFgColor)))) {
                 if (DEBUG) {
                     Slog.d(TAG, "makeSplashScreenContentView: choose fg icon");
                 }
                 // Reference AdaptiveIcon description, outer is 108 and inner is 72, so we
                 // scale by 192/160 if we only draw adaptiveIcon's foreground.
                 final float noBgScale =
-                        foreIconTester.nonTransparentRatio() < ENLARGE_FOREGROUND_ICON_THRESHOLD
+                        iconColor.mFgNonTransparentRatio < ENLARGE_FOREGROUND_ICON_THRESHOLD
                                 ? NO_BACKGROUND_SCALE : 1f;
                 // Using AdaptiveIconDrawable here can help keep the shape consistent with the
                 // current settings.
@@ -689,6 +674,150 @@ public class SplashscreenContentDrawer {
         }
     }
 
+    /** Cache the result of {@link DrawableColorTester} to reduce expensive calculation. */
+    @VisibleForTesting
+    static class ColorCache extends BroadcastReceiver {
+        /**
+         * The color may be different according to resource id and configuration (e.g. night mode),
+         * so this allows to cache more than one color per package.
+         */
+        private static final int CACHE_SIZE = 2;
+
+        /** The computed colors of packages. */
+        private final ArrayMap<String, Colors> mColorMap = new ArrayMap<>();
+
+        private static class Colors {
+            final WindowColor[] mWindowColors = new WindowColor[CACHE_SIZE];
+            final IconColor[] mIconColors = new IconColor[CACHE_SIZE];
+        }
+
+        private static class Cache {
+            /** The hash used to check whether this cache is hit. */
+            final int mHash;
+
+            /** The number of times this cache has been reused. */
+            int mReuseCount;
+
+            Cache(int hash) {
+                mHash = hash;
+            }
+        }
+
+        static class WindowColor extends Cache {
+            final int mBgColor;
+
+            WindowColor(int hash, int bgColor) {
+                super(hash);
+                mBgColor = bgColor;
+            }
+        }
+
+        static class IconColor extends Cache {
+            final int mFgColor;
+            final int mBgColor;
+            final boolean mIsBgComplex;
+            final boolean mIsBgGrayscale;
+            final float mFgNonTransparentRatio;
+
+            IconColor(int hash, int fgColor, int bgColor, boolean isBgComplex,
+                    boolean isBgGrayscale, float fgNonTransparentRatio) {
+                super(hash);
+                mFgColor = fgColor;
+                mBgColor = bgColor;
+                mIsBgComplex = isBgComplex;
+                mIsBgGrayscale = isBgGrayscale;
+                mFgNonTransparentRatio = fgNonTransparentRatio;
+            }
+        }
+
+        ColorCache(Context context, Handler handler) {
+            // This includes reinstall and uninstall.
+            final IntentFilter filter = new IntentFilter(Intent.ACTION_PACKAGE_REMOVED);
+            filter.addDataScheme(IntentFilter.SCHEME_PACKAGE);
+            context.registerReceiverAsUser(this, UserHandle.ALL, filter,
+                    null /* broadcastPermission */, handler);
+        }
+
+        @Override
+        public void onReceive(Context context, Intent intent) {
+            final Uri packageUri = intent.getData();
+            if (packageUri != null) {
+                mColorMap.remove(packageUri.getEncodedSchemeSpecificPart());
+            }
+        }
+
+        /**
+         * Gets the existing cache if the hash matches. If null is returned, the caller can use
+         * outLeastUsedIndex to put the new cache.
+         */
+        private static <T extends Cache> T getCache(T[] caches, int hash, int[] outLeastUsedIndex) {
+            int minReuseCount = Integer.MAX_VALUE;
+            for (int i = 0; i < CACHE_SIZE; i++) {
+                final T cache = caches[i];
+                if (cache == null) {
+                    // Empty slot has the highest priority to put new cache.
+                    minReuseCount = -1;
+                    outLeastUsedIndex[0] = i;
+                    continue;
+                }
+                if (cache.mHash == hash) {
+                    cache.mReuseCount++;
+                    return cache;
+                }
+                if (cache.mReuseCount < minReuseCount) {
+                    minReuseCount = cache.mReuseCount;
+                    outLeastUsedIndex[0] = i;
+                }
+            }
+            return null;
+        }
+
+        @NonNull WindowColor getWindowColor(String packageName, int configHash, int windowBgColor,
+                int windowBgResId, IntSupplier windowBgColorSupplier) {
+            Colors colors = mColorMap.get(packageName);
+            int hash = 31 * configHash + windowBgColor;
+            hash = 31 * hash + windowBgResId;
+            final int[] leastUsedIndex = { 0 };
+            if (colors != null) {
+                final WindowColor windowColor = getCache(colors.mWindowColors, hash,
+                        leastUsedIndex);
+                if (windowColor != null) {
+                    return windowColor;
+                }
+            } else {
+                colors = new Colors();
+                mColorMap.put(packageName, colors);
+            }
+            final WindowColor windowColor = new WindowColor(hash, windowBgColorSupplier.getAsInt());
+            colors.mWindowColors[leastUsedIndex[0]] = windowColor;
+            return windowColor;
+        }
+
+        @NonNull IconColor getIconColor(String packageName, int configHash, int iconResId,
+                Supplier<DrawableColorTester> fgColorTesterSupplier,
+                Supplier<DrawableColorTester> bgColorTesterSupplier) {
+            Colors colors = mColorMap.get(packageName);
+            final int hash = configHash * 31 + iconResId;
+            final int[] leastUsedIndex = { 0 };
+            if (colors != null) {
+                final IconColor iconColor = getCache(colors.mIconColors, hash, leastUsedIndex);
+                if (iconColor != null) {
+                    return iconColor;
+                }
+            } else {
+                colors = new Colors();
+                mColorMap.put(packageName, colors);
+            }
+            final DrawableColorTester fgTester = fgColorTesterSupplier.get();
+            final DrawableColorTester bgTester = bgColorTesterSupplier.get();
+            final IconColor iconColor = new IconColor(hash, fgTester.getDominateColor(),
+                    bgTester.getDominateColor(), bgTester.isComplexColor(), bgTester.isGrayscale(),
+                    fgTester.nonTransparentRatio());
+            colors.mIconColors[leastUsedIndex[0]] = iconColor;
+            return iconColor;
+        }
+    }
+
     /**
      * Create and play the default exit animation for splash screen view.
      */
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 ff91d827b280..8463da6a0ecb 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
@@ -17,6 +17,7 @@
 package com.android.wm.shell.startingsurface;
 
 import static android.content.Context.CONTEXT_RESTRICTED;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.view.Choreographer.CALLBACK_INSETS_ANIMATION;
 import static android.view.Display.DEFAULT_DISPLAY;
 
@@ -34,6 +35,7 @@ import android.graphics.Rect;
 import android.hardware.display.DisplayManager;
 import android.os.IBinder;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.os.UserHandle;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -49,8 +51,10 @@ import android.window.StartingWindowInfo;
 import android.window.TaskSnapshot;
 
 import com.android.internal.R;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.common.annotations.ShellSplashscreenThread;
 
 import java.util.function.Supplier;
 
@@ -90,6 +94,7 @@ import java.util.function.Supplier;
  * => WM#addView -> .. waiting for Choreographer#doFrame -> relayout -> draw -> (draw the Paint
  * directly).
  */
+@ShellSplashscreenThread
 public class StartingSurfaceDrawer {
     static final String TAG = StartingSurfaceDrawer.class.getSimpleName();
     static final boolean DEBUG_SPLASH_SCREEN = StartingWindowController.DEBUG_SPLASH_SCREEN;
@@ -98,7 +103,8 @@ public class StartingSurfaceDrawer {
     private final Context mContext;
     private final DisplayManager mDisplayManager;
     private final ShellExecutor mSplashScreenExecutor;
-    private final SplashscreenContentDrawer mSplashscreenContentDrawer;
+    @VisibleForTesting
+    final SplashscreenContentDrawer mSplashscreenContentDrawer;
     private Choreographer mChoreographer;
 
     private static final boolean DEBUG_ENABLE_REVEAL_ANIMATION =
@@ -276,6 +282,7 @@ public class StartingSurfaceDrawer {
         final SplashScreenViewSupplier viewSupplier = new SplashScreenViewSupplier();
         final FrameLayout rootLayout = new FrameLayout(context);
         final Runnable setViewSynchronized = () -> {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "addSplashScreenView");
             // waiting for setContentView before relayoutWindow
             SplashScreenView contentView = viewSupplier.get();
             final StartingWindowRecord record = mStartingWindowRecords.get(taskId);
@@ -294,13 +301,14 @@ public class StartingSurfaceDrawer {
                 }
                 record.setSplashScreenView(contentView);
             }
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         };
         mSplashscreenContentDrawer.createContentView(context, emptyView, activityInfo, taskId,
                 viewSupplier::setView);
 
         try {
             final WindowManager wm = context.getSystemService(WindowManager.class);
-            if (postAddWindow(taskId, appToken, rootLayout, wm, params)) {
+            if (addWindow(taskId, appToken, rootLayout, wm, params)) {
                 // 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,
@@ -393,10 +401,11 @@ public class StartingSurfaceDrawer {
         ActivityTaskManager.getInstance().onSplashScreenViewCopyFinished(taskId, parcelable);
     }
 
-    protected boolean postAddWindow(int taskId, IBinder appToken, View view, WindowManager wm,
+    protected boolean addWindow(int taskId, IBinder appToken, View view, WindowManager wm,
             WindowManager.LayoutParams params) {
         boolean shouldSaveView = true;
         try {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "addRootView");
             wm.addView(view, params);
         } catch (WindowManager.BadTokenException e) {
             // ignore
@@ -404,6 +413,7 @@ public class StartingSurfaceDrawer {
                     + e.getMessage());
             shouldSaveView = false;
         } finally {
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
             if (view != null && view.getParent() == null) {
                 Slog.w(TAG, "view not successfully added to wm, removing view");
                 wm.removeViewImmediate(view);
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 4e3e133f85b0..903e63ad6554 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
@@ -18,25 +18,27 @@ package com.android.wm.shell.startingsurface;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManager;
+import android.content.BroadcastReceiver;
 import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.graphics.Rect;
+import android.net.Uri;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
+import android.os.UserHandle;
 import android.testing.TestableContext;
 import android.view.SurfaceControl;
 import android.view.View;
@@ -58,6 +60,8 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.function.IntSupplier;
+
 /**
  * Tests for the starting surface drawer.
  */
@@ -71,6 +75,9 @@ public class StartingSurfaceDrawerTests {
     @Mock
     private TransactionPool mTransactionPool;
 
+    private final Handler mTestHandler = new Handler(Looper.getMainLooper());
+    private final TestableContext mTestContext = new TestContext(
+            InstrumentationRegistry.getInstrumentation().getTargetContext());
     TestStartingSurfaceDrawer mStartingSurfaceDrawer;
 
     static final class TestStartingSurfaceDrawer extends StartingSurfaceDrawer{
@@ -83,7 +90,7 @@ public class StartingSurfaceDrawerTests {
         }
 
         @Override
-        protected boolean postAddWindow(int taskId, IBinder appToken,
+        protected boolean addWindow(int taskId, IBinder appToken,
                 View view, WindowManager wm, WindowManager.LayoutParams params) {
             // listen for addView
             mAddWindowForTask = taskId;
@@ -101,45 +108,50 @@ public class StartingSurfaceDrawerTests {
         }
     }
 
+    private static class TestContext extends TestableContext {
+        TestContext(Context context) {
+            super(context);
+        }
+
+        @Override
+        public Context createPackageContextAsUser(String packageName, int flags, UserHandle user)
+                throws PackageManager.NameNotFoundException {
+            return this;
+        }
+
+        @Override
+        public Intent registerReceiverAsUser(BroadcastReceiver receiver, UserHandle user,
+                IntentFilter filter, String broadcastPermission, Handler scheduler) {
+            return null;
+        }
+    }
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        final TestableContext context = new TestableContext(
-                InstrumentationRegistry.getInstrumentation().getTargetContext(), null);
-        final WindowManager realWindowManager = context.getSystemService(WindowManager.class);
+        final WindowManager realWindowManager = mTestContext.getSystemService(WindowManager.class);
         final WindowMetrics metrics = realWindowManager.getMaximumWindowMetrics();
-        context.addMockSystemService(WindowManager.class, mMockWindowManager);
-
-        spyOn(context);
-        spyOn(realWindowManager);
-        try {
-            doReturn(context).when(context)
-                    .createPackageContextAsUser(anyString(), anyInt(), any());
-        } catch (PackageManager.NameNotFoundException e) {
-            //
-        }
+        mTestContext.addMockSystemService(WindowManager.class, mMockWindowManager);
+
         doReturn(metrics).when(mMockWindowManager).getMaximumWindowMetrics();
         doNothing().when(mMockWindowManager).addView(any(), any());
 
-        final HandlerExecutor testExecutor =
-                new HandlerExecutor(new Handler(Looper.getMainLooper()));
-        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(context, testExecutor,
-                mTransactionPool));
+        mStartingSurfaceDrawer = spy(new TestStartingSurfaceDrawer(mTestContext,
+                new HandlerExecutor(mTestHandler), mTransactionPool));
     }
 
     @Test
     public void testAddSplashScreenSurface() {
         final int taskId = 1;
-        final Handler mainLoop = new Handler(Looper.getMainLooper());
         final StartingWindowInfo windowInfo =
                 createWindowInfo(taskId, android.R.style.Theme);
         mStartingSurfaceDrawer.addSplashScreenStartingWindow(windowInfo, mBinder, false);
-        waitHandlerIdle(mainLoop);
-        verify(mStartingSurfaceDrawer).postAddWindow(eq(taskId), eq(mBinder), any(), any(), any());
+        waitHandlerIdle(mTestHandler);
+        verify(mStartingSurfaceDrawer).addWindow(eq(taskId), eq(mBinder), any(), any(), any());
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, taskId);
 
         mStartingSurfaceDrawer.removeStartingWindow(windowInfo.taskInfo.taskId, null, null, false);
-        waitHandlerIdle(mainLoop);
+        waitHandlerIdle(mTestHandler);
         verify(mStartingSurfaceDrawer).removeWindowSynced(eq(taskId), any(), any(), eq(false));
         assertEquals(mStartingSurfaceDrawer.mAddWindowForTask, 0);
     }
@@ -147,15 +159,45 @@ public class StartingSurfaceDrawerTests {
     @Test
     public void testFallbackDefaultTheme() {
         final int taskId = 1;
-        final Handler mainLoop = new Handler(Looper.getMainLooper());
         final StartingWindowInfo windowInfo =
                 createWindowInfo(taskId, 0);
         mStartingSurfaceDrawer.addSplashScreenStartingWindow(windowInfo, mBinder, false);
-        waitHandlerIdle(mainLoop);
-        verify(mStartingSurfaceDrawer).postAddWindow(eq(taskId), eq(mBinder), any(), any(), any());
+        waitHandlerIdle(mTestHandler);
+        verify(mStartingSurfaceDrawer).addWindow(eq(taskId), eq(mBinder), any(), any(), any());
         assertNotEquals(mStartingSurfaceDrawer.mViewThemeResId, 0);
     }
 
+    @Test
+    public void testColorCache() {
+        final String packageName = mTestContext.getPackageName();
+        final int configHash = 1;
+        final int windowBgColor = 0xff000000;
+        final int windowBgResId = 1;
+        final IntSupplier windowBgColorSupplier = () -> windowBgColor;
+        final SplashscreenContentDrawer.ColorCache colorCache =
+                mStartingSurfaceDrawer.mSplashscreenContentDrawer.mColorCache;
+        final SplashscreenContentDrawer.ColorCache.WindowColor windowColor1 =
+                colorCache.getWindowColor(packageName, configHash, windowBgColor, windowBgResId,
+                        windowBgColorSupplier);
+        assertEquals(windowBgColor, windowColor1.mBgColor);
+        assertEquals(0, windowColor1.mReuseCount);
+
+        final SplashscreenContentDrawer.ColorCache.WindowColor windowColor2 =
+                colorCache.getWindowColor(packageName, configHash, windowBgColor, windowBgResId,
+                        windowBgColorSupplier);
+        assertEquals(windowColor1, windowColor2);
+        assertEquals(1, windowColor1.mReuseCount);
+
+        final Intent packageRemoved = new Intent(Intent.ACTION_PACKAGE_REMOVED);
+        packageRemoved.setData(Uri.parse("package:" + packageName));
+        colorCache.onReceive(mTestContext, packageRemoved);
+
+        final SplashscreenContentDrawer.ColorCache.WindowColor windowColor3 =
+                colorCache.getWindowColor(packageName, configHash, windowBgColor, windowBgResId,
+                        windowBgColorSupplier);
+        assertEquals(0, windowColor3.mReuseCount);
+    }
+
     private StartingWindowInfo createWindowInfo(int taskId, int themeResId) {
         StartingWindowInfo windowInfo = new StartingWindowInfo();
         final ActivityInfo info = new ActivityInfo();