Merge "Force calculator to launch in fullscreen for quickswitch tests" into main
diff --git a/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java b/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java
index 0d0f700..556d29c 100644
--- a/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java
+++ b/go/quickstep/src/com/android/launcher3/model/AppShareabilityManager.java
@@ -35,6 +35,7 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.launcher3.model.AppShareabilityDatabase.ShareabilityDao;
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.SafeCloseable;
 
 import java.lang.annotation.Retention;
 import java.util.ArrayList;
@@ -47,7 +48,7 @@
  * Each app's status is retrieved from the Play Store's API. Statuses are cached in order
  * to limit extraneous calls to that API (which can be time-consuming).
  */
-public class AppShareabilityManager {
+public class AppShareabilityManager implements SafeCloseable {
     @Retention(SOURCE)
     @IntDef({
         ShareabilityStatus.UNKNOWN,
@@ -194,6 +195,11 @@
         }
     }
 
+    @Override
+    public void close() {
+        mDatabase.close();
+    }
+
     /**
      * Provides a testable instance of this class
      * This instance allows database queries on the main thread
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index fb14f9e..65a49bd 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -326,8 +326,12 @@
         super.destroy();
         mActive = false;
         StatsLogCompatManager.LOGS_CONSUMER.remove(mAppEventProducer);
-        if (mIsPrimaryInstance) {
-            mStatsManager.clearPullAtomCallback(SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT);
+        if (mIsPrimaryInstance && mStatsManager != null) {
+            try {
+                mStatsManager.clearPullAtomCallback(SysUiStatsLog.LAUNCHER_LAYOUT_SNAPSHOT);
+            } catch (RuntimeException e) {
+                Log.e(TAG, "Failed to unregister snapshot logging callback with StatsManager", e);
+            }
         }
         destroyPredictors();
     }
diff --git a/quickstep/src/com/android/launcher3/model/WellbeingModel.java b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
index c345d6e..a7c9652 100644
--- a/quickstep/src/com/android/launcher3/model/WellbeingModel.java
+++ b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
@@ -32,7 +32,6 @@
 import android.os.Bundle;
 import android.os.DeadObjectException;
 import android.os.Handler;
-import android.os.Looper;
 import android.os.Process;
 import android.os.UserHandle;
 import android.text.TextUtils;
@@ -48,9 +47,10 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.popup.RemoteActionShortcut;
 import com.android.launcher3.popup.SystemShortcut;
-import com.android.launcher3.util.BgObjectWithLooper;
+import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SimpleBroadcastReceiver;
 import com.android.launcher3.views.ActivityContext;
 
@@ -61,7 +61,7 @@
 /**
  * Data model for digital wellbeing status of apps.
  */
-public final class WellbeingModel extends BgObjectWithLooper {
+public final class WellbeingModel implements SafeCloseable {
     private static final String TAG = "WellbeingModel";
     private static final int[] RETRY_TIMES_MS = {5000, 15000, 30000};
     private static final boolean DEBUG = false;
@@ -81,8 +81,12 @@
     private final Context mContext;
     private final String mWellbeingProviderPkg;
 
-    private Handler mWorkerHandler;
-    private ContentObserver mContentObserver;
+    private final Handler mWorkerHandler;
+    private final ContentObserver mContentObserver;
+    private final SimpleBroadcastReceiver mWellbeingAppChangeReceiver =
+            new SimpleBroadcastReceiver(t -> restartObserver());
+    private final SimpleBroadcastReceiver mAppAddRemoveReceiver =
+            new SimpleBroadcastReceiver(this::onAppPackageChanged);
 
     private final Object mModelLock = new Object();
     // Maps the action Id to the corresponding RemoteAction
@@ -94,16 +98,23 @@
     private WellbeingModel(final Context context) {
         mContext = context;
         mWellbeingProviderPkg = mContext.getString(R.string.wellbeing_provider_pkg);
-        initializeInBackground("WellbeingHandler");
+        mWorkerHandler = new Handler(TextUtils.isEmpty(mWellbeingProviderPkg)
+                ? Executors.UI_HELPER_EXECUTOR.getLooper()
+                : Executors.getPackageExecutor(mWellbeingProviderPkg).getLooper());
+
+        mContentObserver = new ContentObserver(mWorkerHandler) {
+            @Override
+            public void onChange(boolean selfChange, Uri uri) {
+                updateAllPackages();
+            }
+        };
+        mWorkerHandler.post(this::initializeInBackground);
     }
 
-    @Override
-    protected void onInitialized(Looper looper) {
-        mWorkerHandler = new Handler(looper);
-        mContentObserver = newContentObserver(mWorkerHandler, this::onWellbeingUriChanged);
+    private void initializeInBackground() {
         if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
             mContext.registerReceiver(
-                    new SimpleBroadcastReceiver(t -> restartObserver()),
+                    mWellbeingAppChangeReceiver,
                     getPackageFilter(mWellbeingProviderPkg,
                             Intent.ACTION_PACKAGE_ADDED, Intent.ACTION_PACKAGE_CHANGED,
                             Intent.ACTION_PACKAGE_REMOVED, Intent.ACTION_PACKAGE_DATA_CLEARED,
@@ -113,17 +124,21 @@
             IntentFilter filter = new IntentFilter(Intent.ACTION_PACKAGE_ADDED);
             filter.addAction(Intent.ACTION_PACKAGE_REMOVED);
             filter.addDataScheme("package");
-            mContext.registerReceiver(new SimpleBroadcastReceiver(this::onAppPackageChanged),
-                    filter, null, mWorkerHandler);
+            mContext.registerReceiver(mAppAddRemoveReceiver, filter, null, mWorkerHandler);
 
             restartObserver();
         }
     }
 
-    @WorkerThread
-    private void onWellbeingUriChanged(Uri uri) {
-        Preconditions.assertNonUiThread();
-        updateAllPackages();
+    @Override
+    public void close() {
+        if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
+            mWorkerHandler.post(() -> {
+                mWellbeingAppChangeReceiver.unregisterReceiverSafely(mContext);
+                mAppAddRemoveReceiver.unregisterReceiverSafely(mContext);
+                mContext.getContentResolver().unregisterContentObserver(mContentObserver);
+            });
+        }
     }
 
     public void setInTest(boolean inTest) {
diff --git a/quickstep/src/com/android/quickstep/RotationTouchHelper.java b/quickstep/src/com/android/quickstep/RotationTouchHelper.java
index bf50d70..3380291 100644
--- a/quickstep/src/com/android/quickstep/RotationTouchHelper.java
+++ b/quickstep/src/com/android/quickstep/RotationTouchHelper.java
@@ -39,6 +39,7 @@
 import com.android.launcher3.util.DisplayController.Info;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.NavigationMode;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.TaskStackChangeListener;
@@ -50,7 +51,7 @@
 /**
  * Helper class for transforming touch events
  */
-public class RotationTouchHelper implements DisplayInfoChangeListener {
+public class RotationTouchHelper implements DisplayInfoChangeListener, SafeCloseable {
 
     public static final MainThreadInitializedObject<RotationTouchHelper> INSTANCE =
             new MainThreadInitializedObject<>(RotationTouchHelper::new);
@@ -197,6 +198,11 @@
         mOnDestroyActions.add(action);
     }
 
+    @Override
+    public void close() {
+        destroy();
+    }
+
     /**
      * Cleans up all the registered listeners and receivers.
      */
diff --git a/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
index f474796..29a57fc 100644
--- a/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/SimpleOrientationTouchTransformer.java
@@ -24,22 +24,30 @@
 
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.SafeCloseable;
 
 public class SimpleOrientationTouchTransformer implements
-        DisplayController.DisplayInfoChangeListener {
+        DisplayController.DisplayInfoChangeListener, SafeCloseable {
 
     public static final MainThreadInitializedObject<SimpleOrientationTouchTransformer> INSTANCE =
             new MainThreadInitializedObject<>(SimpleOrientationTouchTransformer::new);
 
+    private final Context mContext;
     private OrientationRectF mOrientationRectF;
 
     public SimpleOrientationTouchTransformer(Context context) {
+        mContext = context;
         DisplayController.INSTANCE.get(context).addChangeListener(this);
         onDisplayInfoChanged(context, DisplayController.INSTANCE.get(context).getInfo(),
                 CHANGE_ALL);
     }
 
     @Override
+    public void close() {
+        DisplayController.INSTANCE.get(mContext).removeChangeListener(this);
+    }
+
+    @Override
     public void onDisplayInfoChanged(Context context, DisplayController.Info info, int flags) {
         if ((flags & (CHANGE_ROTATION | CHANGE_ACTIVE_SCREEN)) == 0) {
             return;
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index edf05b5..ab9840f 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -67,6 +67,7 @@
 import com.android.internal.view.AppearanceRegion;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.quickstep.util.ActiveGestureLog;
 import com.android.quickstep.util.AssistUtils;
 import com.android.quickstep.util.unfold.ProxyUnfoldTransitionProvider;
@@ -110,7 +111,7 @@
 /**
  * Holds the reference to SystemUI.
  */
-public class SystemUiProxy implements ISystemUiProxy, NavHandle {
+public class SystemUiProxy implements ISystemUiProxy, NavHandle, SafeCloseable {
     private static final String TAG = SystemUiProxy.class.getSimpleName();
 
     public static final MainThreadInitializedObject<SystemUiProxy> INSTANCE =
@@ -201,6 +202,9 @@
     }
 
     @Override
+    public void close() { }
+
+    @Override
     public void onBackPressed() {
         if (mSystemUiProxy != null) {
             try {
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index dec8a12..9d899fc 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -59,6 +59,7 @@
     public static final boolean SHELL_TRANSITIONS_ROTATION = ENABLE_SHELL_TRANSITIONS
             && SystemProperties.getBoolean("persist.wm.debug.shell_transit_rotate", false);
 
+    private final Context mCtx;
     private RecentsAnimationController mController;
     private RecentsAnimationCallbacks mCallbacks;
     private RecentsAnimationTargets mTargets;
@@ -66,7 +67,6 @@
     private GestureState mLastGestureState;
     private RemoteAnimationTarget[] mLastAppearedTaskTargets;
     private Runnable mLiveTileCleanUpHandler;
-    private Context mCtx;
 
     private boolean mRecentsAnimationStartPending = false;
     private boolean mShouldIgnoreMotionEvents = false;
@@ -329,7 +329,7 @@
                 options.setTransientLaunch();
             }
             options.setSourceInfo(ActivityOptions.SourceInfo.TYPE_RECENTS_ANIMATION, eventTime);
-            mRecentsAnimationStartPending = SystemUiProxy.INSTANCE.getNoCreate()
+            mRecentsAnimationStartPending = SystemUiProxy.INSTANCE.get(mCtx)
                     .startRecentsActivity(intent, options, mCallbacks);
             if (enableHandleDelayedGestureCallbacks()) {
                 ActiveGestureLog.INSTANCE.addLog(new ActiveGestureLog.CompoundString(
diff --git a/quickstep/src/com/android/quickstep/TaskUtils.java b/quickstep/src/com/android/quickstep/TaskUtils.java
index 80a449b..63e536a 100644
--- a/quickstep/src/com/android/quickstep/TaskUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskUtils.java
@@ -70,7 +70,7 @@
             return "";
         }
         UserHandle user = UserHandle.of(userId);
-        ApplicationInfo applicationInfo = new PackageManagerHelper(context)
+        ApplicationInfo applicationInfo = PackageManagerHelper.INSTANCE.get(context)
                 .getApplicationInfo(packageName, user, 0);
         if (applicationInfo == null) {
             Log.e(TAG, "Failed to get title for userId=" + userId + ", packageName=" + packageName);
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index a2a6dde..3a6b804 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -34,6 +34,7 @@
 import androidx.annotation.UiThread;
 
 import com.android.launcher3.util.MainThreadInitializedObject;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.SplitStageInfo;
 import com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
@@ -57,7 +58,8 @@
  * This class tracked the top-most task and  some 'approximate' task history to allow faster
  * system state estimation during touch interaction
  */
-public class TopTaskTracker extends ISplitScreenListener.Stub implements TaskStackChangeListener {
+public class TopTaskTracker extends ISplitScreenListener.Stub
+        implements TaskStackChangeListener, SafeCloseable {
 
     public static MainThreadInitializedObject<TopTaskTracker> INSTANCE =
             new MainThreadInitializedObject<>(TopTaskTracker::new);
@@ -67,12 +69,13 @@
     // Ordered list with first item being the most recent task.
     private final LinkedList<RunningTaskInfo> mOrderedTaskList = new LinkedList<>();
 
-
+    private final Context mContext;
     private final SplitStageInfo mMainStagePosition = new SplitStageInfo();
     private final SplitStageInfo mSideStagePosition = new SplitStageInfo();
     private int mPinnedTaskId = INVALID_TASK_ID;
 
     private TopTaskTracker(Context context) {
+        mContext = context;
         mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
         mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
 
@@ -81,6 +84,12 @@
     }
 
     @Override
+    public void close() {
+        TaskStackChangeListeners.getInstance().unregisterTaskStackListener(this);
+        SystemUiProxy.INSTANCE.get(mContext).unregisterSplitScreenListener(this);
+    }
+
+    @Override
     public void onTaskRemoved(int taskId) {
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskId);
     }
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index a09e027..ed633df 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -347,7 +347,6 @@
                         event.getId() + "";
                 Log.d(TAG, name);
             }
-            LauncherAppState appState = LauncherAppState.getInstanceNoCreate();
 
             if (mSlice == null && mSliceItem != null) {
                 mSlice = LauncherAtom.Slice.newBuilder().setUri(
@@ -369,15 +368,10 @@
                 return;
             }
 
-            if (mItemInfo.container < 0 || appState == null) {
-                // Write log on the model thread so that logs do not go out of order
-                // (for eg: drop comes after drag)
-                Executors.MODEL_EXECUTOR.execute(
-                        () -> write(event, applyOverwrites(mItemInfo.buildProto())));
-            } else {
+            if (mItemInfo.container < 0 || !LauncherAppState.INSTANCE.executeIfCreated(app -> {
                 // Item is inside a collection, fetch collection info in a BG thread
                 // and then write to StatsLog.
-                appState.getModel().enqueueModelUpdateTask(
+                app.getModel().enqueueModelUpdateTask(
                         new BaseModelUpdateTask() {
                             @Override
                             public void execute(@NonNull final LauncherAppState app,
@@ -388,6 +382,11 @@
                                 write(event, applyOverwrites(mItemInfo.buildProto(collectionInfo)));
                             }
                         });
+            })) {
+                // Write log on the model thread so that logs do not go out of order
+                // (for eg: drop comes after drag)
+                Executors.MODEL_EXECUTOR.execute(
+                        () -> write(event, applyOverwrites(mItemInfo.buildProto())));
             }
         }
 
diff --git a/quickstep/src/com/android/quickstep/util/AppPairsController.java b/quickstep/src/com/android/quickstep/util/AppPairsController.java
index 2b4d280..a82031a 100644
--- a/quickstep/src/com/android/quickstep/util/AppPairsController.java
+++ b/quickstep/src/com/android/quickstep/util/AppPairsController.java
@@ -316,7 +316,7 @@
                 itemInfos.stream().map(ItemInfo::getComponentKey).toList();
 
         // Use TopTaskTracker to find the currently running app (or apps)
-        TopTaskTracker topTaskTracker = getTopTaskTracker(context);
+        TopTaskTracker topTaskTracker = getTopTaskTracker();
 
         // getRunningSplitTasksIds() will return a pair of ids if we are currently running a
         // split pair, or an empty array with zero length if we are running a single app.
@@ -489,7 +489,7 @@
      * Gets the TopTaskTracker, which is a cached record of the top running Task.
      */
     @VisibleForTesting
-    public TopTaskTracker getTopTaskTracker(Context context) {
-        return TopTaskTracker.INSTANCE.get(context);
+    public TopTaskTracker getTopTaskTracker() {
+        return TopTaskTracker.INSTANCE.get(mContext);
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
index b6e6bf7..2396512 100644
--- a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
@@ -63,7 +63,7 @@
             SplitSelectStateController controller) {
         mLauncher = launcher;
         mController = controller;
-        mIconCache = LauncherAppState.getInstanceNoCreate().getIconCache();
+        mIconCache = LauncherAppState.getInstance(launcher).getIconCache();
         mHalfDividerSize = mLauncher.getResources().getDimensionPixelSize(
                 R.dimen.multi_window_task_divider_size) / 2;
     }
diff --git a/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java b/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java
index 89d8cc4..e80d2a6 100644
--- a/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java
+++ b/quickstep/src/com/android/quickstep/util/TaskRemovedDuringLaunchListener.java
@@ -22,6 +22,8 @@
 import static com.android.launcher3.BaseActivity.EVENT_RESUMED;
 import static com.android.launcher3.BaseActivity.EVENT_STOPPED;
 
+import android.content.Context;
+
 import androidx.annotation.NonNull;
 
 import com.android.quickstep.RecentsModel;
@@ -45,6 +47,12 @@
     private final Runnable mUnregisterCallback = this::unregister;
     private final Runnable mResumeCallback = this::checkTaskLaunchFailed;
 
+    private final Context mContext;
+
+    public TaskRemovedDuringLaunchListener(Context context) {
+        mContext = context;
+    }
+
     /**
      * Registers a failure listener callback if it detects a scenario in which an app launch
      * failed before the transition finished.
@@ -88,7 +96,7 @@
         if (mLaunchedTaskId != INVALID_TASK_ID) {
             final int launchedTaskId = mLaunchedTaskId;
             final Runnable taskLaunchFailedCallback = mTaskLaunchFailedCallback;
-            RecentsModel.INSTANCE.getNoCreate().isTaskRemoved(mLaunchedTaskId, (taskRemoved) -> {
+            RecentsModel.INSTANCE.get(mContext).isTaskRemoved(mLaunchedTaskId, (taskRemoved) -> {
                 if (taskRemoved) {
                     ActiveGestureLog.INSTANCE.addLog(
                             new ActiveGestureLog.CompoundString("Launch failed, task (id=")
diff --git a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
index 8fa5375..82ba30b 100644
--- a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
+++ b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
@@ -19,7 +19,7 @@
 import static android.provider.Settings.ACTION_APP_USAGE_SETTINGS;
 
 import static com.android.launcher3.Utilities.prefixTextWithIcon;
-import static com.android.launcher3.util.Executors.THREAD_POOL_EXECUTOR;
+import static com.android.launcher3.util.Executors.ORDERED_BG_EXECUTOR;
 
 import android.app.ActivityOptions;
 import android.content.ActivityNotFoundException;
@@ -140,7 +140,7 @@
     public void initialize(Task task) {
         mAppUsageLimitTimeMs = mAppRemainingTimeMs = -1;
         mTask = task;
-        THREAD_POOL_EXECUTOR.execute(() -> {
+        ORDERED_BG_EXECUTOR.execute(() -> {
                     AppUsageLimit usageLimit = null;
                     try {
                         usageLimit = mLauncherApps.getAppUsageLimit(
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index f88a498..a033243 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -91,6 +91,7 @@
 import com.android.launcher3.util.CancellableTask;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.RunnableList;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.SplitPositionOption;
 import com.android.launcher3.util.TraceHelper;
@@ -116,6 +117,8 @@
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.QuickStepContract;
 
+import kotlin.Unit;
+
 import java.lang.annotation.Retention;
 import java.util.Arrays;
 import java.util.Collections;
@@ -124,8 +127,6 @@
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
-import kotlin.Unit;
-
 /**
  * A task in the Recents view.
  */
@@ -912,8 +913,8 @@
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
 
-            final TaskRemovedDuringLaunchListener
-                    failureListener = new TaskRemovedDuringLaunchListener();
+            TaskRemovedDuringLaunchListener failureListener = new TaskRemovedDuringLaunchListener(
+                    getContext().getApplicationContext());
             if (isQuickswitch) {
                 // We only listen for failures to launch in quickswitch because the during this
                 // gesture launcher is in the background state, vs other launches which are in
@@ -940,12 +941,8 @@
             // Indicate success once the system has indicated that the transition has started
             ActivityOptions opts = ActivityOptions.makeCustomTaskAnimation(getContext(), 0, 0,
                     MAIN_EXECUTOR.getHandler(),
-                    elapsedRealTime -> {
-                        callback.accept(true);
-                    },
-                    elapsedRealTime -> {
-                        failureListener.onTransitionFinished();
-                    });
+                    elapsedRealTime -> callback.accept(true),
+                    elapsedRealTime -> failureListener.onTransitionFinished());
             opts.setLaunchDisplayId(
                     getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
             if (isQuickswitch) {
@@ -1840,7 +1837,7 @@
     /**
      * We update and subsequently draw these in {@link #setFullscreenProgress(float)}.
      */
-    public static class FullscreenDrawParams {
+    public static class FullscreenDrawParams implements SafeCloseable {
 
         private float mCornerRadius;
         private float mWindowCornerRadius;
@@ -1875,6 +1872,9 @@
                     Utilities.mapRange(fullscreenProgress, mCornerRadius, mWindowCornerRadius)
                             / parentScale / taskViewScale;
         }
+
+        @Override
+        public void close() { }
     }
 
     public class TaskIdAttributeContainer {
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
index 9fa4b79..72cfd92 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
@@ -44,7 +44,6 @@
 import com.android.launcher3.util.window.CachedDisplayInfo;
 import com.android.launcher3.util.window.WindowManagerProxy;
 import com.android.quickstep.FallbackActivityInterface;
-import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.util.SurfaceTransaction.MockProperties;
 
 import org.hamcrest.Description;
@@ -160,7 +159,6 @@
         void verifyNoTransforms() {
             LauncherModelHelper helper = new LauncherModelHelper();
             try {
-                helper.sandboxContext.allow(SystemUiProxy.INSTANCE);
                 int rotation = mDisplaySize.x > mDisplaySize.y
                         ? Surface.ROTATION_90 : Surface.ROTATION_0;
                 CachedDisplayInfo cdi = new CachedDisplayInfo(mDisplaySize, rotation);
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
index f0683f9..ec245ee 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
@@ -15,6 +15,8 @@
  */
 package com.android.quickstep;
 
+import static androidx.test.InstrumentationRegistry.getTargetContext;
+
 import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
 import static com.android.quickstep.TaplTestsTaskbar.TaskbarMode.PERSISTENT;
 import static com.android.quickstep.TaplTestsTaskbar.TaskbarMode.TRANSIENT;
@@ -53,7 +55,7 @@
 
     @Override
     public void setUp() throws Exception {
-        mTaskbarWasInTransientMode = isTaskbarInTransientMode(mTargetContext);
+        mTaskbarWasInTransientMode = isTaskbarInTransientMode(getTargetContext());
         setTaskbarMode(mLauncher, isTaskbarTestModeTransient());
         super.setUp();
     }
diff --git a/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt
index adaf7ff..ece67af 100644
--- a/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/AppPairsControllerTest.kt
@@ -43,6 +43,7 @@
 import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.never
 import org.mockito.kotlin.spy
@@ -100,8 +101,7 @@
         // Stub methods on appPairsController so that they return mocks
         spyAppPairsController = spy(appPairsController)
         whenever(mockAppPairIcon.context).thenReturn(mockTaskbarActivityContext)
-        whenever(spyAppPairsController.getTopTaskTracker(mockTaskbarActivityContext))
-            .thenReturn(mockTopTaskTracker)
+        doReturn(mockTopTaskTracker).whenever(spyAppPairsController).topTaskTracker
         whenever(mockTopTaskTracker.getCachedTopTask(any())).thenReturn(mockCachedTaskInfo)
         whenever(mockTask1.getKey()).thenReturn(mockTaskKey1)
         whenever(mockTask2.getKey()).thenReturn(mockTaskKey2)
diff --git a/res/drawable/ic_private_profile_app_scroller_badge.xml b/res/drawable/ic_private_profile_app_scroller_badge.xml
index b52a277..ede42b9 100644
--- a/res/drawable/ic_private_profile_app_scroller_badge.xml
+++ b/res/drawable/ic_private_profile_app_scroller_badge.xml
@@ -21,8 +21,8 @@
     <path
         android:pathData="M16.0007 2.66602L5.33398 6.66602V14.786C5.33398 21.5194 9.88065 27.7993 16.0007 29.3327C22.1207 27.7993 26.6673 21.5194 26.6673 14.786V6.66602L16.0007 2.66602ZM20.0007 19.9993V22.666H17.334V23.9993H14.6673V17.1193C12.7473 16.546 11.334 14.786 11.334 12.666C11.334 10.0927 13.4273 7.99935 16.0007 7.99935C18.574 7.99935 20.6673 10.0927 20.6673 12.666C20.6673 14.7727 19.254 16.546 17.334 17.1193V19.9993H20.0007Z"
         android:fillType="evenOdd"
-        android:fillColor="@android:color/white" />
+        android:fillColor="?android:attr/textColorPrimaryInverse" />
     <path
         android:pathData="M16 14.666C17.1046 14.666 18 13.7706 18 12.666C18 11.5614 17.1046 10.666 16 10.666C14.8954 10.666 14 11.5614 14 12.666C14 13.7706 14.8954 14.666 16 14.666Z"
-        android:fillColor="@android:color/white" />
+        android:fillColor="?android:attr/textColorPrimaryInverse" />
 </vector>
diff --git a/res/layout/private_space_header.xml b/res/layout/private_space_header.xml
index 65f1004..84554e4 100644
--- a/res/layout/private_space_header.xml
+++ b/res/layout/private_space_header.xml
@@ -40,7 +40,6 @@
             android:layout_width="@dimen/ps_header_image_height"
             android:layout_height="@dimen/ps_header_image_height"
             android:background="@drawable/ps_settings_background"
-            android:layout_marginEnd="@dimen/ps_header_settings_icon_margin_end"
             android:src="@drawable/ic_ps_settings"
             android:contentDescription="@string/ps_container_settings" />
         <LinearLayout
@@ -49,7 +48,7 @@
             android:layout_height="@dimen/ps_header_image_height"
             android:background="@drawable/ps_lock_background"
             android:gravity="center_vertical"
-            android:layout_marginEnd="@dimen/ps_header_layout_margin"
+            android:layout_marginEnd="@dimen/ps_lock_button_margin_end"
             android:contentDescription="@string/ps_container_lock_unlock_button">
             <ImageView
                 android:id="@+id/lock_icon"
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 5ff9902..5dd3569 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -496,7 +496,7 @@
     <dimen name="ps_header_image_height">48dp</dimen>
     <dimen name="ps_header_text_height">24dp</dimen>
     <dimen name="ps_header_layout_margin">16dp</dimen>
-    <dimen name="ps_header_settings_icon_margin_end">4dp</dimen>
+    <dimen name="ps_lock_button_margin_end">12dp</dimen>
     <dimen name="ps_header_text_size">16sp</dimen>
     <dimen name="ps_button_height">40dp</dimen>
     <dimen name="ps_button_width">40dp</dimen>
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 98cb84e..f405b93 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -59,6 +59,7 @@
 import com.android.launcher3.util.LockedUserState;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.Partner;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
 
@@ -75,7 +76,7 @@
 import java.util.Objects;
 import java.util.stream.Collectors;
 
-public class InvariantDeviceProfile {
+public class InvariantDeviceProfile implements SafeCloseable {
 
     public static final String TAG = "IDP";
     // We do not need any synchronization for this variable as its only written on UI thread.
@@ -229,9 +230,8 @@
         if (!newGridName.equals(gridName)) {
             LauncherPrefs.get(context).put(GRID_NAME, newGridName);
         }
-        LockedUserState.get(context).runOnUserUnlocked(() -> {
-            new DeviceGridState(this).writeToPrefs(context);
-        });
+        LockedUserState.get(context).runOnUserUnlocked(() ->
+            new DeviceGridState(this).writeToPrefs(context));
 
         DisplayController.INSTANCE.get(context).setPriorityListener(
                 (displayContext, info, flags) -> {
@@ -295,6 +295,11 @@
         initGrid(context, myInfo, result, deviceType);
     }
 
+    @Override
+    public void close() {
+        DisplayController.INSTANCE.executeIfCreated(dc -> dc.setPriorityListener(null));
+    }
+
     /**
      * Reinitialize the current grid after a restore, where some grids might now be disabled.
      */
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 50a597d..d2633e0 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -78,14 +78,10 @@
 
     private final RunnableList mOnTerminateCallback = new RunnableList();
 
-    public static LauncherAppState getInstance(final Context context) {
+    public static LauncherAppState getInstance(Context context) {
         return INSTANCE.get(context);
     }
 
-    public static LauncherAppState getInstanceNoCreate() {
-        return INSTANCE.getNoCreate();
-    }
-
     public Context getContext() {
         return mContext;
     }
diff --git a/src/com/android/launcher3/LauncherModel.java b/src/com/android/launcher3/LauncherModel.java
index 7fdfd72..9b0e0ec 100644
--- a/src/com/android/launcher3/LauncherModel.java
+++ b/src/com/android/launcher3/LauncherModel.java
@@ -447,8 +447,8 @@
                 IconCache iconCache = app.getIconCache();
                 final IntSet removedIds = new IntSet();
                 HashSet<WorkspaceItemInfo> archivedWorkspaceItemsToCacheRefresh = new HashSet<>();
-                boolean isAppArchived = new PackageManagerHelper(
-                        mApp.getContext()).isAppArchivedForUser(packageName, user);
+                boolean isAppArchived = PackageManagerHelper.INSTANCE.get(mApp.getContext())
+                        .isAppArchivedForUser(packageName, user);
                 synchronized (dataModel) {
                     if (isAppArchived) {
                         // Remove package icon cache entry for archived app in case of a session
diff --git a/src/com/android/launcher3/LauncherPrefs.kt b/src/com/android/launcher3/LauncherPrefs.kt
index 80a8cea..b503739 100644
--- a/src/com/android/launcher3/LauncherPrefs.kt
+++ b/src/com/android/launcher3/LauncherPrefs.kt
@@ -30,6 +30,7 @@
 import com.android.launcher3.states.RotationHelper
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.MainThreadInitializedObject
+import com.android.launcher3.util.SafeCloseable
 import com.android.launcher3.util.Themes
 
 /**
@@ -37,7 +38,7 @@
  *
  * TODO(b/262721340): Replace all direct SharedPreference refs with LauncherPrefs / Item methods.
  */
-class LauncherPrefs(private val encryptedContext: Context) {
+class LauncherPrefs(private val encryptedContext: Context) : SafeCloseable {
     private val deviceProtectedStorageContext =
         encryptedContext.createDeviceProtectedStorageContext()
 
@@ -242,6 +243,8 @@
         }
     }
 
+    override fun close() {}
+
     companion object {
         @VisibleForTesting const val BOOT_AWARE_PREFS_KEY = "boot_aware_prefs"
 
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 4e0ba62..6e2d357 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -48,11 +48,11 @@
      */
     @Override
     public void dump(FileDescriptor fd, PrintWriter writer, String[] args) {
-        LauncherAppState appState = LauncherAppState.getInstanceNoCreate();
-        if (appState == null || !appState.getModel().isModelLoaded()) {
-            return;
-        }
-        appState.getModel().dumpState("", fd, writer, args);
+        LauncherAppState.INSTANCE.executeIfCreated(appState -> {
+            if (appState.getModel().isModelLoaded()) {
+                appState.getModel().dumpState("", fd, writer, args);
+            }
+        });
     }
 
     @Override
diff --git a/src/com/android/launcher3/SecondaryDropTarget.java b/src/com/android/launcher3/SecondaryDropTarget.java
index 1362586..0a4fb73 100644
--- a/src/com/android/launcher3/SecondaryDropTarget.java
+++ b/src/com/android/launcher3/SecondaryDropTarget.java
@@ -362,7 +362,7 @@
 
         public void onLauncherResume() {
             // We use MATCH_UNINSTALLED_PACKAGES as the app can be on SD card as well.
-            if (new PackageManagerHelper(mContext).getApplicationInfo(mPackageName,
+            if (PackageManagerHelper.INSTANCE.get(mContext).getApplicationInfo(mPackageName,
                     mDragObject.dragInfo.user, PackageManager.MATCH_UNINSTALLED_PACKAGES) == null) {
                 mDragObject.dragSource = mOriginal;
                 mOriginal.onDropCompleted(SecondaryDropTarget.this, mDragObject, true);
diff --git a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
index 60df7c5..a810331 100644
--- a/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
+++ b/src/com/android/launcher3/allapps/AlphabeticalAppsList.java
@@ -22,7 +22,6 @@
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_PRIVATE_SPACE_USER_INSTALLED_APPS_COUNT;
 
 import android.content.Context;
-import android.graphics.drawable.Drawable;
 import android.text.Spannable;
 import android.text.SpannableString;
 import android.text.style.ImageSpan;
@@ -362,7 +361,7 @@
         // Split of private space apps into user-installed and system apps.
         Map<Boolean, List<AppInfo>> split = mPrivateApps.stream()
                 .collect(Collectors.partitioningBy(mPrivateProviderManager
-                                .splitIntoUserInstalledAndSystemApps()));
+                                .splitIntoUserInstalledAndSystemApps(mActivityContext)));
 
         // TODO(b/329688630): switch to the pulled LayoutStaticSnapshot atom
         mActivityContext
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 38fe138..7e975df 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -80,9 +80,7 @@
 import com.android.launcher3.views.RecyclerViewFastScroller;
 
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import java.util.function.Predicate;
 
 /**
@@ -105,7 +103,6 @@
             }
         }
     };
-    private Set<String> mPreInstalledSystemPackages = new HashSet<>();
     private Intent mAppInstallerIntent = new Intent();
     private PrivateAppsSectionDecorator mPrivateAppsSectionDecorator;
     private boolean mPrivateSpaceSettingsAvailable;
@@ -264,8 +261,6 @@
         ApiWrapper apiWrapper = ApiWrapper.INSTANCE.get(appContext);
         UserHandle profileUser = getProfileUser();
         if (profileUser != null) {
-            mPreInstalledSystemPackages = new HashSet<>(
-                    apiWrapper.getPreInstalledSystemPackages(profileUser));
             mAppInstallerIntent = apiWrapper
                     .getAppMarketActivityIntent(BuildConfig.APPLICATION_ID, profileUser);
         }
@@ -349,10 +344,12 @@
      * Splits private apps into user installed and system apps.
      * When the list of system apps is empty, all apps are treated as system.
      */
-    public Predicate<AppInfo> splitIntoUserInstalledAndSystemApps() {
-        return appInfo -> !mPreInstalledSystemPackages.isEmpty()
+    public Predicate<AppInfo> splitIntoUserInstalledAndSystemApps(Context context) {
+        List<String> preInstallApps = UserCache.getInstance(context)
+                .getPreInstallApps(getProfileUser());
+        return appInfo -> !preInstallApps.isEmpty()
                 && (appInfo.componentName == null
-                || !(mPreInstalledSystemPackages.contains(appInfo.componentName.getPackageName())));
+                || !(preInstallApps.contains(appInfo.componentName.getPackageName())));
     }
 
     /** Add Private Space Header view elements based upon {@link UserProfileState} */
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index 05fdcef..ec0a222 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -162,7 +162,7 @@
             finish();
             return;
         }
-        ApplicationInfo info = new PackageManagerHelper(this)
+        ApplicationInfo info = PackageManagerHelper.INSTANCE.get(this)
                 .getApplicationInfo(targetApp.packageName, targetApp.user, 0);
         if (info == null) {
             finish();
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 6bed9dc..50d2b43 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -66,7 +66,6 @@
 import com.android.launcher3.InsettableFrameLayout;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
@@ -88,14 +87,11 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
-import com.android.launcher3.util.PluginManagerWrapper;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.util.window.WindowManagerProxy;
 import com.android.launcher3.views.ActivityContext;
@@ -104,7 +100,6 @@
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 import com.android.launcher3.widget.LocalColorExtractor;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
 import com.android.launcher3.widget.util.WidgetSizes;
 
 import java.util.ArrayList;
@@ -136,13 +131,10 @@
                 new ConcurrentLinkedQueue<>();
 
         public PreviewContext(Context base, InvariantDeviceProfile idp) {
-            super(base, UserCache.INSTANCE, InstallSessionHelper.INSTANCE, LauncherPrefs.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    CustomWidgetManager.INSTANCE, PluginManagerWrapper.INSTANCE,
-                    WindowManagerProxy.INSTANCE, DisplayController.INSTANCE);
+            super(base);
             mIdp = idp;
-            mObjectMap.put(InvariantDeviceProfile.INSTANCE, idp);
-            mObjectMap.put(LauncherAppState.INSTANCE,
+            putObject(InvariantDeviceProfile.INSTANCE, idp);
+            putObject(LauncherAppState.INSTANCE,
                     new LauncherAppState(this, null /* iconCacheFileName */));
         }
 
diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index ce563b7..3fa6da4 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.model;
 
+import android.content.Context;
 import android.content.Intent;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
@@ -85,6 +86,7 @@
 
         final ArrayList<ItemInfo> addedItemsFinal = new ArrayList<>();
         final IntArray addedWorkspaceScreensFinal = new IntArray();
+        final Context context = app.getContext();
 
         synchronized (dataModel) {
             IntArray workspaceScreens = dataModel.collectWorkspaceScreens();
@@ -99,7 +101,7 @@
                     }
 
                     // b/139663018 Short-circuit this logic if the icon is a system app
-                    if (PackageManagerHelper.isSystemApp(app.getContext(),
+                    if (PackageManagerHelper.isSystemApp(context,
                             Objects.requireNonNull(item.getIntent()))) {
                         continue;
                     }
@@ -112,7 +114,7 @@
 
                 if (item.itemType == LauncherSettings.Favorites.ITEM_TYPE_APPLICATION) {
                     if (item instanceof WorkspaceItemFactory) {
-                        item = ((WorkspaceItemFactory) item).makeWorkspaceItem(app.getContext());
+                        item = ((WorkspaceItemFactory) item).makeWorkspaceItem(context);
                     }
                 }
                 if (item != null) {
@@ -121,8 +123,8 @@
             }
 
             InstallSessionHelper packageInstaller =
-                    InstallSessionHelper.INSTANCE.get(app.getContext());
-            LauncherApps launcherApps = app.getContext().getSystemService(LauncherApps.class);
+                    InstallSessionHelper.INSTANCE.get(context);
+            LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
 
             for (ItemInfo item : filteredItems) {
                 // Find appropriate space for the item.
@@ -135,7 +137,7 @@
                         || item instanceof LauncherAppWidgetInfo) {
                     itemInfo = item;
                 } else if (item instanceof WorkspaceItemFactory) {
-                    itemInfo = ((WorkspaceItemFactory) item).makeWorkspaceItem(app.getContext());
+                    itemInfo = ((WorkspaceItemFactory) item).makeWorkspaceItem(context);
                 } else {
                     throw new RuntimeException("Unexpected info type");
                 }
diff --git a/src/com/android/launcher3/model/AllAppsList.java b/src/com/android/launcher3/model/AllAppsList.java
index a1a05f4..39c1243 100644
--- a/src/com/android/launcher3/model/AllAppsList.java
+++ b/src/com/android/launcher3/model/AllAppsList.java
@@ -169,7 +169,7 @@
     public AppInfo addPromiseApp(
             Context context, PackageInstallInfo installInfo, boolean loadIcon) {
         // only if not yet installed
-        if (new PackageManagerHelper(context)
+        if (PackageManagerHelper.INSTANCE.get(context)
                 .isAppInstalled(installInfo.packageName, installInfo.user)) {
             return null;
         }
diff --git a/src/com/android/launcher3/model/ItemInstallQueue.java b/src/com/android/launcher3/model/ItemInstallQueue.java
index 90aba2a..551c2d8 100644
--- a/src/com/android/launcher3/model/ItemInstallQueue.java
+++ b/src/com/android/launcher3/model/ItemInstallQueue.java
@@ -45,7 +45,6 @@
 import com.android.launcher3.Launcher;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
@@ -55,6 +54,7 @@
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.PersistedItemArray;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
 import java.util.HashSet;
@@ -65,7 +65,7 @@
 /**
  * Class to maintain a queue of pending items to be added to the workspace.
  */
-public class ItemInstallQueue {
+public class ItemInstallQueue implements SafeCloseable {
 
     private static final String LOG = "ItemInstallQueue";
 
@@ -99,6 +99,9 @@
         mContext = context;
     }
 
+    @Override
+    public void close() {}
+
     @WorkerThread
     private void ensureQueueLoaded() {
         Preconditions.assertWorkerThread();
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index a742c75..77bc32e 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -138,6 +138,7 @@
     private final LauncherApps mLauncherApps;
     private final UserManager mUserManager;
     private final UserCache mUserCache;
+    private final PackageManagerHelper mPmHelper;
 
     private final InstallSessionHelper mSessionHelper;
     private final IconCache mIconCache;
@@ -170,6 +171,7 @@
         mLauncherApps = mApp.getContext().getSystemService(LauncherApps.class);
         mUserManager = mApp.getContext().getSystemService(UserManager.class);
         mUserCache = UserCache.INSTANCE.get(mApp.getContext());
+        mPmHelper = PackageManagerHelper.INSTANCE.get(mApp.getContext());
         mSessionHelper = InstallSessionHelper.INSTANCE.get(mApp.getContext());
         mIconCache = mApp.getIconCache();
         mUserManagerState = userManagerState;
@@ -407,7 +409,6 @@
             @Nullable LoaderMemoryLogger memoryLogger,
             @Nullable LauncherRestoreEventLogger restoreEventLogger) {
         final Context context = mApp.getContext();
-        final PackageManagerHelper pmHelper = new PackageManagerHelper(context);
         final boolean isSdCardReady = Utilities.isBootCompleted();
         final WidgetInflater widgetInflater = new WidgetInflater(context);
 
@@ -447,7 +448,7 @@
                         mUserCache, mUserManagerState, mLauncherApps, mPendingPackages,
                         mShortcutKeyToPinnedShortcuts, mApp, mBgDataModel,
                         mWidgetProvidersMap, installingPkgs, isSdCardReady,
-                        widgetInflater, pmHelper, iconRequestInfos, unlockedUsers,
+                        widgetInflater, mPmHelper, iconRequestInfos, unlockedUsers,
                         allDeepShortcuts);
 
                 while (!mStopped && c.moveToNext()) {
diff --git a/src/com/android/launcher3/model/SdCardAvailableReceiver.java b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
index 8cfa3aa..5293316 100644
--- a/src/com/android/launcher3/model/SdCardAvailableReceiver.java
+++ b/src/com/android/launcher3/model/SdCardAvailableReceiver.java
@@ -52,7 +52,7 @@
     @Override
     public void onReceive(Context context, Intent intent) {
         final LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
-        final PackageManagerHelper pmHelper = new PackageManagerHelper(context);
+        final PackageManagerHelper pmHelper = PackageManagerHelper.INSTANCE.get(context);
         for (PackageUserKey puk : mPackages) {
             UserHandle user = puk.mUser;
 
diff --git a/src/com/android/launcher3/model/ShortcutsChangedTask.java b/src/com/android/launcher3/model/ShortcutsChangedTask.java
index 59dd1b1..1cb5215 100644
--- a/src/com/android/launcher3/model/ShortcutsChangedTask.java
+++ b/src/com/android/launcher3/model/ShortcutsChangedTask.java
@@ -78,7 +78,7 @@
 
         if (!matchingWorkspaceItems.isEmpty()) {
             if (mShortcuts.isEmpty()) {
-                PackageManagerHelper packageManagerHelper = new PackageManagerHelper(
+                PackageManagerHelper packageManagerHelper = PackageManagerHelper.INSTANCE.get(
                         app.getContext());
                 // Verify that the app is indeed installed.
                 if (!packageManagerHelper.isAppInstalled(mPackageName, mUser)
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index 8c3efd7..72e85c7 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -427,12 +427,10 @@
     @NonNull
     protected LauncherAtom.ItemInfo.Builder getDefaultItemInfoBuilder() {
         LauncherAtom.ItemInfo.Builder itemBuilder = LauncherAtom.ItemInfo.newBuilder();
-        UserIconInfo info = getUserInfo();
-        itemBuilder.setIsWork(info != null && info.isWork());
-        itemBuilder.setUserType(getUserType(info));
-        SettingsCache settingsCache = SettingsCache.INSTANCE.getNoCreate();
-        boolean isKidsMode = settingsCache != null && settingsCache.getValue(NAV_BAR_KIDS_MODE, 0);
-        itemBuilder.setIsKidsMode(isKidsMode);
+        SettingsCache.INSTANCE.executeIfCreated(cache ->
+                itemBuilder.setIsKidsMode(cache.getValue(NAV_BAR_KIDS_MODE, 0)));
+        UserCache.INSTANCE.executeIfCreated(cache ->
+                itemBuilder.setUserType(getUserType(cache.getUserInfo(user))));
         itemBuilder.setRank(rank);
         return itemBuilder;
     }
@@ -526,15 +524,6 @@
         this.title = title;
     }
 
-    private UserIconInfo getUserInfo() {
-        UserCache userCache = UserCache.INSTANCE.getNoCreate();
-        if (userCache == null) {
-            return null;
-        }
-
-        return userCache.getUserInfo(user);
-    }
-
     private int getUserType(UserIconInfo info) {
         if (info == null) {
             return SysUiStatsLog.LAUNCHER_UICHANGED__USER_TYPE__TYPE_UNKNOWN;
diff --git a/src/com/android/launcher3/pm/InstallSessionHelper.java b/src/com/android/launcher3/pm/InstallSessionHelper.java
index 4a3318e..e66f496 100644
--- a/src/com/android/launcher3/pm/InstallSessionHelper.java
+++ b/src/com/android/launcher3/pm/InstallSessionHelper.java
@@ -32,7 +32,6 @@
 import com.android.launcher3.Flags;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.SessionCommitReceiver;
-import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.ItemInstallQueue;
 import com.android.launcher3.util.IntArray;
@@ -41,6 +40,7 @@
 import com.android.launcher3.util.PackageManagerHelper;
 import com.android.launcher3.util.PackageUserKey;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.SafeCloseable;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -52,7 +52,7 @@
  * Utility class to tracking install sessions
  */
 @SuppressWarnings("NewApi")
-public class InstallSessionHelper {
+public class InstallSessionHelper implements SafeCloseable {
 
     @NonNull
     private static final String LOG = "InstallSessionHelper";
@@ -89,6 +89,9 @@
         mLauncherApps = context.getSystemService(LauncherApps.class);
     }
 
+    @Override
+    public void close() { }
+
     @WorkerThread
     @NonNull
     private IntSet getPromiseIconIds() {
@@ -168,7 +171,7 @@
         synchronized (mSessionVerifiedMap) {
             if (!mSessionVerifiedMap.containsKey(pkg)) {
                 boolean hasSystemFlag = DEBUG || mAppContext.getPackageName().equals(pkg)
-                        || new PackageManagerHelper(mAppContext)
+                        || PackageManagerHelper.INSTANCE.get(mAppContext)
                                 .getApplicationInfo(pkg, user, ApplicationInfo.FLAG_SYSTEM) != null;
                 mSessionVerifiedMap.put(pkg, hasSystemFlag);
             }
@@ -242,7 +245,7 @@
                 && sessionInfo.getInstallReason() == PackageManager.INSTALL_REASON_USER
                 && sessionInfo.getAppIcon() != null
                 && !TextUtils.isEmpty(sessionInfo.getAppLabel())
-                && !new PackageManagerHelper(mAppContext).isAppInstalled(
+                && !PackageManagerHelper.INSTANCE.get(mAppContext).isAppInstalled(
                         sessionInfo.getAppPackageName(), getUserHandle(sessionInfo));
     }
 
diff --git a/src/com/android/launcher3/pm/UserCache.java b/src/com/android/launcher3/pm/UserCache.java
index 185e0b5..b7b557d 100644
--- a/src/com/android/launcher3/pm/UserCache.java
+++ b/src/com/android/launcher3/pm/UserCache.java
@@ -24,6 +24,7 @@
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.util.ArrayMap;
 
 import androidx.annotation.AnyThread;
 import androidx.annotation.NonNull;
@@ -81,6 +82,9 @@
     @NonNull
     private Map<UserHandle, UserIconInfo> mUserToSerialMap;
 
+    @NonNull
+    private Map<UserHandle, List<String>> mUserToPreInstallAppMap;
+
     private UserCache(Context context) {
         mContext = context;
         mUserToSerialMap = Collections.emptyMap();
@@ -120,6 +124,20 @@
     @WorkerThread
     private void updateCache() {
         mUserToSerialMap = ApiWrapper.INSTANCE.get(mContext).queryAllUsers();
+        mUserToPreInstallAppMap = fetchPreInstallApps();
+    }
+
+    @WorkerThread
+    private Map<UserHandle, List<String>> fetchPreInstallApps() {
+        Map<UserHandle, List<String>> userToPreInstallApp = new ArrayMap<>();
+        mUserToSerialMap.forEach((userHandle, userIconInfo) -> {
+            // Fetch only for private profile, as other profiles have no usages yet.
+            List<String> preInstallApp = userIconInfo.isPrivate()
+                    ? ApiWrapper.INSTANCE.get(mContext).getPreInstalledSystemPackages(userHandle)
+                    : new ArrayList<>();
+            userToPreInstallApp.put(userHandle, preInstallApp);
+        });
+        return userToPreInstallApp;
     }
 
     /**
@@ -172,6 +190,15 @@
     }
 
     /**
+     * Returns the pre-installed apps for a user.
+     */
+    @NonNull
+    public List<String> getPreInstallApps(UserHandle user) {
+        List<String> preInstallApp = mUserToPreInstallAppMap.get(user);
+        return preInstallApp == null ? new ArrayList<>() : preInstallApp;
+    }
+
+    /**
      * Get a non-themed {@link UserBadgeDrawable} based on the provided {@link UserHandle}.
      */
     @Nullable
diff --git a/src/com/android/launcher3/popup/SystemShortcut.java b/src/com/android/launcher3/popup/SystemShortcut.java
index f56d732..6005573 100644
--- a/src/com/android/launcher3/popup/SystemShortcut.java
+++ b/src/com/android/launcher3/popup/SystemShortcut.java
@@ -181,8 +181,8 @@
         public void onClick(View view) {
             dismissTaskMenuView();
             Rect sourceBounds = Utilities.getViewBounds(view);
-            new PackageManagerHelper(view.getContext()).startDetailsActivityForInfo(
-                    mItemInfo, sourceBounds, ActivityOptions.makeBasic().toBundle());
+            PackageManagerHelper.startDetailsActivityForInfo(view.getContext(), mItemInfo,
+                    sourceBounds, ActivityOptions.makeBasic().toBundle());
             mTarget.getStatsLogManager().logger().withItemInfo(mItemInfo)
                     .log(LAUNCHER_SYSTEM_SHORTCUT_APP_INFO_TAP);
         }
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index d5f1e18..a4ff29f 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -526,10 +526,7 @@
         }
 
         logFavoritesTable(controller.getDb(), "launcher db after remap widget ids", null, null);
-        LauncherAppState app = LauncherAppState.getInstanceNoCreate();
-        if (app != null) {
-            app.getModel().forceReload();
-        }
+        LauncherAppState.INSTANCE.executeIfCreated(app -> app.getModel().forceReload());
     }
 
     private static void logDatabaseWidgetInfo(ModelDbController controller) {
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index d5c9b9f..db2a6e0 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -94,13 +94,10 @@
 
     protected Context mContext;
     protected DeviceProfile mDeviceProfile;
-    protected LauncherAppState mLauncherAppState;
 
     public void init(Context context) {
         mContext = context;
-        mDeviceProfile = InvariantDeviceProfile.INSTANCE.
-                get(context).getDeviceProfile(context);
-        mLauncherAppState = LauncherAppState.getInstanceNoCreate();
+        mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(context).getDeviceProfile(context);
         if (sActivityLifecycleCallbacks == null) {
             sActivityLifecycleCallbacks = new ActivityLifecycleCallbacksAdapter() {
                 @Override
diff --git a/src/com/android/launcher3/util/BgObjectWithLooper.java b/src/com/android/launcher3/util/BgObjectWithLooper.java
deleted file mode 100644
index adc3c7d..0000000
--- a/src/com/android/launcher3/util/BgObjectWithLooper.java
+++ /dev/null
@@ -1,63 +0,0 @@
-/*
- * Copyright (C) 2020 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.launcher3.util;
-
-import android.database.ContentObserver;
-import android.net.Uri;
-import android.os.Handler;
-import android.os.Looper;
-
-import androidx.annotation.WorkerThread;
-
-import java.util.function.Consumer;
-
-/**
- * Utility class to define an object which does most of it's processing on a
- * dedicated background thread.
- */
-public abstract class BgObjectWithLooper {
-
-    /**
-     * Start initialization of the object
-     */
-    public final void initializeInBackground(String threadName) {
-        new Thread(this::runOnThread, threadName).start();
-    }
-
-    private void runOnThread() {
-        Looper.prepare();
-        onInitialized(Looper.myLooper());
-        Looper.loop();
-    }
-
-    /**
-     * Called on the background thread to handle initialization
-     */
-    @WorkerThread
-    protected abstract void onInitialized(Looper looper);
-
-    /**
-     * Helper method to create a content provider
-     */
-    protected static ContentObserver newContentObserver(Handler handler, Consumer<Uri> command) {
-        return new ContentObserver(handler) {
-            @Override
-            public void onChange(boolean selfChange, Uri uri) {
-                command.accept(uri);
-            }
-        };
-    }
-}
diff --git a/src/com/android/launcher3/util/DynamicResource.java b/src/com/android/launcher3/util/DynamicResource.java
index 1008ebb..fbdb5c2 100644
--- a/src/com/android/launcher3/util/DynamicResource.java
+++ b/src/com/android/launcher3/util/DynamicResource.java
@@ -33,7 +33,8 @@
  *
  * To allow customization for a particular resource, add them to dynamic_resources.xml
  */
-public class DynamicResource implements ResourceProvider, PluginListener<ResourceProvider> {
+public class DynamicResource implements
+        ResourceProvider, PluginListener<ResourceProvider>, SafeCloseable {
 
     private static final MainThreadInitializedObject<DynamicResource> INSTANCE =
             new MainThreadInitializedObject<>(DynamicResource::new);
@@ -48,6 +49,11 @@
     }
 
     @Override
+    public void close() {
+        PluginManagerWrapper.INSTANCE.get(mContext).removePluginListener(this);
+    }
+
+    @Override
     public int getInt(@IntegerRes int resId) {
         return mContext.getResources().getInteger(resId);
     }
diff --git a/src/com/android/launcher3/util/MainThreadInitializedObject.java b/src/com/android/launcher3/util/MainThreadInitializedObject.java
index b966d8e..1a0f9a0 100644
--- a/src/com/android/launcher3/util/MainThreadInitializedObject.java
+++ b/src/com/android/launcher3/util/MainThreadInitializedObject.java
@@ -28,17 +28,15 @@
 import com.android.launcher3.util.ResourceBasedOverride.Overrides;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ExecutionException;
+import java.util.function.Consumer;
 
 /**
  * Utility class for defining singletons which are initiated on main thread.
  */
-public class MainThreadInitializedObject<T> {
+public class MainThreadInitializedObject<T extends SafeCloseable> {
 
     private final ObjectProvider<T> mProvider;
     private T mValue;
@@ -48,14 +46,14 @@
     }
 
     public T get(Context context) {
-        if (context instanceof SandboxContext sc) {
+        Context app = context.getApplicationContext();
+        if (app instanceof SandboxApplication sc) {
             return sc.getObject(this);
         }
 
         if (mValue == null) {
             if (Looper.myLooper() == Looper.getMainLooper()) {
-                mValue = TraceHelper.allowIpcs("main.thread.object",
-                        () -> mProvider.get(context.getApplicationContext()));
+                mValue = TraceHelper.allowIpcs("main.thread.object", () -> mProvider.get(app));
             } else {
                 try {
                     return MAIN_EXECUTOR.submit(() -> get(context)).get();
@@ -67,8 +65,18 @@
         return mValue;
     }
 
-    public T getNoCreate() {
-        return mValue;
+    /**
+     * Executes the callback is the value is already created
+     * @return true if the callback was executed, false otherwise
+     */
+    public boolean executeIfCreated(Consumer<T> callback) {
+        T v = mValue;
+        if (v != null) {
+            callback.accept(v);
+            return true;
+        } else {
+            return false;
+        }
     }
 
     @VisibleForTesting
@@ -79,8 +87,8 @@
     /**
      * Initializes a provider based on resource overrides
      */
-    public static <T extends ResourceBasedOverride> MainThreadInitializedObject<T> forOverride(
-            Class<T> clazz, int resourceId) {
+    public static <T extends ResourceBasedOverride & SafeCloseable> MainThreadInitializedObject<T>
+            forOverride(Class<T> clazz, int resourceId) {
         return new MainThreadInitializedObject<>(c -> Overrides.getObject(clazz, c, resourceId));
     }
 
@@ -89,24 +97,36 @@
         T get(Context context);
     }
 
+    public interface SandboxApplication {
+
+        /**
+         * Find a cached object from mObjectMap if we have already created one. If not, generate
+         * an object using the provider.
+         */
+        <T extends SafeCloseable> T getObject(MainThreadInitializedObject<T> object);
+
+        @UiThread
+        default <T extends SafeCloseable> T createObject(MainThreadInitializedObject<T> object) {
+            return object.mProvider.get((Context) this);
+        }
+    }
+
     /**
      * Abstract Context which allows custom implementations for
      * {@link MainThreadInitializedObject} providers
      */
-    public static class SandboxContext extends ContextWrapper {
+    public static class SandboxContext extends ContextWrapper implements SandboxApplication {
 
         private static final String TAG = "SandboxContext";
 
-        protected final Set<MainThreadInitializedObject> mAllowedObjects;
-        protected final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
-        protected final ArrayList<Object> mOrderedObjects = new ArrayList<>();
+        private final Map<MainThreadInitializedObject, Object> mObjectMap = new HashMap<>();
+        private final ArrayList<SafeCloseable> mOrderedObjects = new ArrayList<>();
 
         private final Object mDestroyLock = new Object();
         private boolean mDestroyed = false;
 
-        public SandboxContext(Context base, MainThreadInitializedObject... allowedObjects) {
+        public SandboxContext(Context base) {
             super(base);
-            mAllowedObjects = new HashSet<>(Arrays.asList(allowedObjects));
         }
 
         @Override
@@ -118,20 +138,14 @@
             synchronized (mDestroyLock) {
                 // Destroy in reverse order
                 for (int i = mOrderedObjects.size() - 1; i >= 0; i--) {
-                    Object o = mOrderedObjects.get(i);
-                    if (o instanceof SafeCloseable) {
-                        ((SafeCloseable) o).close();
-                    }
+                    mOrderedObjects.get(i).close();
                 }
                 mDestroyed = true;
             }
         }
 
-        /**
-         * Find a cached object from mObjectMap if we have already created one. If not, generate
-         * an object using the provider.
-         */
-        protected <T> T getObject(MainThreadInitializedObject<T> object) {
+        @Override
+        public <T extends SafeCloseable> T getObject(MainThreadInitializedObject<T> object) {
             synchronized (mDestroyLock) {
                 if (mDestroyed) {
                     Log.e(TAG, "Static object access with a destroyed context");
@@ -142,12 +156,6 @@
                 }
                 if (Looper.myLooper() == Looper.getMainLooper()) {
                     t = createObject(object);
-                    // Check if we've explicitly allowed the object or if it's a SafeCloseable,
-                    // it will get destroyed in onDestroy()
-                    if (!mAllowedObjects.contains(object) && !(t instanceof SafeCloseable)) {
-                        throw new IllegalStateException("Leaking unknown objects "
-                                + object + "  " + object.mProvider + " " + t);
-                    }
                     mObjectMap.put(object, t);
                     mOrderedObjects.add(t);
                     return t;
@@ -161,17 +169,12 @@
             }
         }
 
-        @UiThread
-        protected <T> T createObject(MainThreadInitializedObject<T> object) {
-            return object.mProvider.get(this);
-        }
-
         /**
          * Put a value into mObjectMap, can be used to put mocked MainThreadInitializedObject
          * instances into SandboxContext.
          */
-        @VisibleForTesting
-        public <T> void putObject(MainThreadInitializedObject<T> object, T value) {
+        public <T extends SafeCloseable> void putObject(
+                MainThreadInitializedObject<T> object, T value) {
             mObjectMap.put(object, value);
         }
     }
diff --git a/src/com/android/launcher3/util/PackageManagerHelper.java b/src/com/android/launcher3/util/PackageManagerHelper.java
index 608bed7..3684f56 100644
--- a/src/com/android/launcher3/util/PackageManagerHelper.java
+++ b/src/com/android/launcher3/util/PackageManagerHelper.java
@@ -56,11 +56,15 @@
 /**
  * Utility methods using package manager
  */
-public class PackageManagerHelper {
+public class PackageManagerHelper implements SafeCloseable{
 
     private static final String TAG = "PackageManagerHelper";
 
     @NonNull
+    public static final MainThreadInitializedObject<PackageManagerHelper> INSTANCE =
+            new MainThreadInitializedObject<>(PackageManagerHelper::new);
+
+    @NonNull
     private final Context mContext;
 
     @NonNull
@@ -75,6 +79,9 @@
         mLauncherApps = Objects.requireNonNull(context.getSystemService(LauncherApps.class));
     }
 
+    @Override
+    public void close() { }
+
     /**
      * Returns true if the app can possibly be on the SDCard. This is just a workaround and doesn't
      * guarantee that the app is on SD card.
@@ -170,10 +177,11 @@
     /**
      * Starts the details activity for {@code info}
      */
-    public void startDetailsActivityForInfo(ItemInfo info, Rect sourceBounds, Bundle opts) {
+    public static void startDetailsActivityForInfo(Context context, ItemInfo info,
+            Rect sourceBounds, Bundle opts) {
         if (info instanceof ItemInfoWithIcon appInfo
                 && (appInfo.runtimeStatusFlags & FLAG_INSTALL_SESSION_ACTIVE) != 0) {
-            mContext.startActivity(ApiWrapper.INSTANCE.get(mContext).getAppMarketActivityIntent(
+            context.startActivity(ApiWrapper.INSTANCE.get(context).getAppMarketActivityIntent(
                     appInfo.getTargetComponent().getPackageName(), Process.myUserHandle()));
             return;
         }
@@ -189,9 +197,10 @@
         }
         if (componentName != null) {
             try {
-                mLauncherApps.startAppDetailsActivity(componentName, info.user, sourceBounds, opts);
+                context.getSystemService(LauncherApps.class).startAppDetailsActivity(componentName,
+                        info.user, sourceBounds, opts);
             } catch (SecurityException | ActivityNotFoundException e) {
-                Toast.makeText(mContext, R.string.activity_not_found, Toast.LENGTH_SHORT).show();
+                Toast.makeText(context, R.string.activity_not_found, Toast.LENGTH_SHORT).show();
                 Log.e(TAG, "Unable to launch settings", e);
             }
         }
diff --git a/src/com/android/launcher3/util/ScreenOnTracker.java b/src/com/android/launcher3/util/ScreenOnTracker.java
index 67530a6..e16e477 100644
--- a/src/com/android/launcher3/util/ScreenOnTracker.java
+++ b/src/com/android/launcher3/util/ScreenOnTracker.java
@@ -27,7 +27,7 @@
 /**
  * Utility class for tracking if the screen is currently on or off
  */
-public class ScreenOnTracker {
+public class ScreenOnTracker implements SafeCloseable {
 
     public static final MainThreadInitializedObject<ScreenOnTracker> INSTANCE =
             new MainThreadInitializedObject<>(ScreenOnTracker::new);
@@ -35,14 +35,21 @@
     private final SimpleBroadcastReceiver mReceiver = new SimpleBroadcastReceiver(this::onReceive);
     private final CopyOnWriteArrayList<ScreenOnListener> mListeners = new CopyOnWriteArrayList<>();
 
+    private final Context mContext;
     private boolean mIsScreenOn;
 
     private ScreenOnTracker(Context context) {
         // Assume that the screen is on to begin with
+        mContext = context;
         mIsScreenOn = true;
         mReceiver.register(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
     }
 
+    @Override
+    public void close() {
+        mReceiver.unregisterReceiverSafely(mContext);
+    }
+
     private void onReceive(Intent intent) {
         String action = intent.getAction();
         if (ACTION_SCREEN_ON.equals(action)) {
diff --git a/src/com/android/launcher3/util/VibratorWrapper.java b/src/com/android/launcher3/util/VibratorWrapper.java
index cd60c1d..b2b4959 100644
--- a/src/com/android/launcher3/util/VibratorWrapper.java
+++ b/src/com/android/launcher3/util/VibratorWrapper.java
@@ -39,7 +39,7 @@
 /**
  * Wrapper around {@link Vibrator} to easily perform haptic feedback where necessary.
  */
-public class VibratorWrapper {
+public class VibratorWrapper implements SafeCloseable {
 
     public static final MainThreadInitializedObject<VibratorWrapper> INSTANCE =
             new MainThreadInitializedObject<>(VibratorWrapper::new);
@@ -77,6 +77,7 @@
     private final Vibrator mVibrator;
     private final boolean mHasVibrator;
 
+    private ContentObserver mHapticFeedbackObserver;
     private boolean mIsHapticFeedbackEnabled;
 
     private VibratorWrapper(Context context) {
@@ -86,14 +87,14 @@
         if (mHasVibrator) {
             final ContentResolver resolver = context.getContentResolver();
             mIsHapticFeedbackEnabled = isHapticFeedbackEnabled(resolver);
-            final ContentObserver observer = new ContentObserver(MAIN_EXECUTOR.getHandler()) {
+            mHapticFeedbackObserver = new ContentObserver(MAIN_EXECUTOR.getHandler()) {
                 @Override
                 public void onChange(boolean selfChange) {
                     mIsHapticFeedbackEnabled = isHapticFeedbackEnabled(resolver);
                 }
             };
             resolver.registerContentObserver(Settings.System.getUriFor(HAPTIC_FEEDBACK_ENABLED),
-                    false /* notifyForDescendants */, observer);
+                    false /* notifyForDescendants */, mHapticFeedbackObserver);
         } else {
             mIsHapticFeedbackEnabled = false;
         }
@@ -126,6 +127,13 @@
         }
     }
 
+    @Override
+    public void close() {
+        if (mHapticFeedbackObserver != null) {
+            mContext.getContentResolver().unregisterContentObserver(mHapticFeedbackObserver);
+        }
+    }
+
     /**
      * This is called when the user swipes to/from all apps. This is meant to be used in between
      * long animation progresses so that it gives a dragging texture effect. For a better
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index 998191e..4b004f3 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -59,6 +59,7 @@
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.NavigationMode;
 import com.android.launcher3.util.ResourceBasedOverride;
+import com.android.launcher3.util.SafeCloseable;
 import com.android.launcher3.util.WindowBounds;
 
 import java.util.ArrayList;
@@ -67,7 +68,7 @@
 /**
  * Utility class for mocking some window manager behaviours
  */
-public class WindowManagerProxy implements ResourceBasedOverride {
+public class WindowManagerProxy implements ResourceBasedOverride, SafeCloseable {
 
     private static final String TAG = "WindowManagerProxy";
     public static final int MIN_TABLET_WIDTH = 600;
@@ -305,12 +306,12 @@
 
         navBarHeightPortrait = isTablet
                 ? (mTaskbarDrawnInProcess
-                        ? 0 : systemRes.getDimensionPixelSize(R.dimen.taskbar_size))
+                        ? 0 : context.getResources().getDimensionPixelSize(R.dimen.taskbar_size))
                 : getDimenByName(systemRes, NAVBAR_HEIGHT);
 
         navBarHeightLandscape = isTablet
                 ? (mTaskbarDrawnInProcess
-                        ? 0 : systemRes.getDimensionPixelSize(R.dimen.taskbar_size))
+                        ? 0 : context.getResources().getDimensionPixelSize(R.dimen.taskbar_size))
                 : (isTabletOrGesture
                         ? getDimenByName(systemRes, NAVBAR_HEIGHT_LANDSCAPE) : 0);
         navbarWidthLandscape = isTabletOrGesture
@@ -474,6 +475,9 @@
                 NavigationMode.THREE_BUTTONS;
     }
 
+    @Override
+    public void close() { }
+
     /**
      * @see DisplayCutout#getSafeInsets
      */
diff --git a/src/com/android/launcher3/widget/BaseWidgetSheet.java b/src/com/android/launcher3/widget/BaseWidgetSheet.java
index e5b5daa..749554c 100644
--- a/src/com/android/launcher3/widget/BaseWidgetSheet.java
+++ b/src/com/android/launcher3/widget/BaseWidgetSheet.java
@@ -141,9 +141,17 @@
 
         if (enableWidgetTapToAdd()) {
             scrollToWidgetCell(wc);
+
             if (mWidgetCellWithAddButton != null) {
-                // If there is a add button currently showing, hide it.
-                mWidgetCellWithAddButton.hideAddButton(/* animate= */ true);
+                if (mWidgetCellWithAddButton.isShowingAddButton()) {
+                    // If there is a add button currently showing, hide it.
+                    mWidgetCellWithAddButton.hideAddButton(/* animate= */ true);
+                } else {
+                    // The last recorded widget cell to show an add button is no longer showing it,
+                    // likely because the widget cell has been recycled or lost focus. If this is
+                    // the cell that has been clicked, we will show it below.
+                    mWidgetCellWithAddButton = null;
+                }
             }
 
             if (mWidgetCellWithAddButton != wc) {
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index ab007fb..ea167d7 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -110,6 +110,7 @@
     private int mSourceContainer = CONTAINER_WIDGETS_TRAY;
 
     private CancellableTask mIconLoadRequest;
+    private boolean mIsShowingAddButton = false;
 
     public WidgetCell(Context context) {
         this(context, null);
@@ -534,6 +535,9 @@
      * @param callback Callback to be set on the button.
      */
     public void showAddButton(View.OnClickListener callback) {
+        if (mIsShowingAddButton) return;
+        mIsShowingAddButton = true;
+
         mWidgetAddButton.setAlpha(0F);
         mWidgetAddButton.setVisibility(VISIBLE);
         mWidgetAddButton.setOnClickListener(callback);
@@ -555,6 +559,9 @@
      * Hide tap to add button.
      */
     public void hideAddButton(boolean animate) {
+        if (!mIsShowingAddButton) return;
+        mIsShowingAddButton = false;
+
         mWidgetAddButton.setOnClickListener(null);
         mWidgetAddButton.animate().cancel();
         mWidgetTextContainer.animate().cancel();
@@ -579,4 +586,8 @@
                 .alpha(1F)
                 .setDuration(ADD_BUTTON_FADE_DURATION_MS);
     }
+
+    public boolean isShowingAddButton() {
+        return mIsShowingAddButton;
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
index c3906a6..80bda22 100644
--- a/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
+++ b/src/com/android/launcher3/widget/picker/WidgetRecommendationCategoryProvider.java
@@ -62,16 +62,17 @@
         // via the overridden WidgetRecommendationCategoryProvider resource.
 
         Preconditions.assertWorkerThread();
-        PackageManagerHelper pmHelper = new PackageManagerHelper(context);
-        if (item.widgetInfo != null && item.widgetInfo.getComponent() != null) {
-            String widgetComponentName = item.widgetInfo.getComponent().getClassName();
-            ApplicationInfo applicationInfo = pmHelper.getApplicationInfo(
-                    item.widgetInfo.getComponent().getPackageName(), item.widgetInfo.getUser(),
-                    0 /* flags */);
-            if (applicationInfo != null) {
-                int predictionCategory = applicationInfo.category;
-                return getCategoryFromApplicationCategory(context, predictionCategory,
-                        widgetComponentName);
+        try (PackageManagerHelper pmHelper = new PackageManagerHelper(context)) {
+            if (item.widgetInfo != null && item.widgetInfo.getComponent() != null) {
+                String widgetComponentName = item.widgetInfo.getComponent().getClassName();
+                ApplicationInfo applicationInfo = pmHelper.getApplicationInfo(
+                        item.widgetInfo.getComponent().getPackageName(), item.widgetInfo.getUser(),
+                        0 /* flags */);
+                if (applicationInfo != null) {
+                    int predictionCategory = applicationInfo.category;
+                    return getCategoryFromApplicationCategory(context, predictionCategory,
+                            widgetComponentName);
+                }
             }
         }
         return null;
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
index 0009a4e..002f496 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
+++ b/tests/multivalentTests/src/com/android/launcher3/util/LauncherModelHelper.java
@@ -52,17 +52,11 @@
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
-import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.model.AllAppsList;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.ItemInstallQueue;
-import com.android.launcher3.pm.InstallSessionHelper;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.testing.TestInformationProvider;
 import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext;
-import com.android.launcher3.util.window.WindowManagerProxy;
-import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.io.ByteArrayOutputStream;
 import java.io.File;
@@ -233,13 +227,7 @@
         private final File mDbDir;
 
         public SandboxModelContext() {
-            super(ApplicationProvider.getApplicationContext(),
-                    UserCache.INSTANCE, InstallSessionHelper.INSTANCE, LauncherPrefs.INSTANCE,
-                    LauncherAppState.INSTANCE, InvariantDeviceProfile.INSTANCE,
-                    DisplayController.INSTANCE, CustomWidgetManager.INSTANCE,
-                    SettingsCache.INSTANCE, PluginManagerWrapper.INSTANCE,
-                    LockedUserState.INSTANCE, WallpaperColorHints.INSTANCE,
-                    ItemInstallQueue.INSTANCE, WindowManagerProxy.INSTANCE);
+            super(ApplicationProvider.getApplicationContext());
 
             // System settings cache content provider. Ensure that they are statically initialized
             Settings.Secure.getString(
@@ -254,18 +242,13 @@
         }
 
         @Override
-        protected <T> T createObject(MainThreadInitializedObject<T> object) {
+        public <T extends SafeCloseable> T createObject(MainThreadInitializedObject<T> object) {
             if (object == LauncherAppState.INSTANCE) {
                 return (T) new LauncherAppState(this, null /* iconCacheFileName */);
             }
             return super.createObject(object);
         }
 
-        public SandboxModelContext allow(MainThreadInitializedObject object) {
-            mAllowedObjects.add(object);
-            return this;
-        }
-
         @Override
         public File getDatabasePath(String name) {
             if (!mDbDir.exists()) {
diff --git a/tests/shared/com/android/launcher3/testing/OWNERS b/tests/shared/com/android/launcher3/testing/OWNERS
new file mode 100644
index 0000000..02e8ebc
--- /dev/null
+++ b/tests/shared/com/android/launcher3/testing/OWNERS
@@ -0,0 +1,4 @@
+vadimt@google.com
+sunnygoyal@google.com
+winsonc@google.com
+hyunyoungs@google.com
diff --git a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index 577334b..aefc2db 100644
--- a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -300,13 +300,7 @@
                 smallestScreenWidthDp = min(screenWidthDp, screenHeightDp)
             }
         val configurationContext = runningContext.createConfigurationContext(config)
-        context =
-            SandboxContext(
-                configurationContext,
-                DisplayController.INSTANCE,
-                WindowManagerProxy.INSTANCE,
-                LauncherPrefs.INSTANCE
-            )
+        context = SandboxContext(configurationContext)
         context.putObject(DisplayController.INSTANCE, displayController)
         context.putObject(WindowManagerProxy.INSTANCE, windowManagerProxy)
         context.putObject(LauncherPrefs.INSTANCE, launcherPrefs)
diff --git a/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java b/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java
index 423ca24..d2238ff 100644
--- a/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java
+++ b/tests/src/com/android/launcher3/allapps/AlphabeticalAppsListTest.java
@@ -98,7 +98,7 @@
         when(mPrivateProfileManager.addPrivateSpaceHeader(any()))
                 .thenAnswer(answer(this::addPrivateSpaceHeader));
         when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps(any()))
                 .thenReturn(iteminfo -> iteminfo.componentName == null
                         || !iteminfo.componentName.getPackageName()
                         .equals("com.android.launcher3.tests.camera"));
@@ -127,7 +127,7 @@
         when(mPrivateProfileManager.addSystemAppsDivider(any()))
                 .thenAnswer(answer(this::addSystemAppsDivider));
         when(mPrivateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps())
+        when(mPrivateProfileManager.splitIntoUserInstalledAndSystemApps(mContext))
                 .thenReturn(iteminfo -> iteminfo.componentName == null
                         || !iteminfo.componentName.getPackageName()
                         .equals("com.android.launcher3.tests.camera"));
diff --git a/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java b/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java
index 9363ce8..512b2ac 100644
--- a/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java
+++ b/tests/src/com/android/launcher3/allapps/PrivateSpaceHeaderViewTest.java
@@ -71,6 +71,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Predicate;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -279,10 +280,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doAnswer(answer(this::addPrivateSpaceHeader)).when(privateProfileManager)
                 .addPrivateSpaceHeader(any());
@@ -316,10 +315,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doAnswer(answer(this::addPrivateSpaceHeader)).when(privateProfileManager)
                 .addPrivateSpaceHeader(any());
@@ -353,10 +350,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doAnswer(answer(this::addPrivateSpaceHeader)).when(privateProfileManager)
                 .addPrivateSpaceHeader(any());
@@ -390,10 +385,8 @@
         when(mAllAppsStore.getApps()).thenReturn(createAppInfoList());
         PrivateProfileManager privateProfileManager = spy(mPrivateProfileManager);
         when(privateProfileManager.getCurrentState()).thenReturn(STATE_ENABLED);
-        when(privateProfileManager.splitIntoUserInstalledAndSystemApps())
-                .thenReturn(iteminfo -> iteminfo.componentName == null
-                        || !iteminfo.componentName.getPackageName()
-                        .equals(CAMERA_PACKAGE_NAME));
+        doReturn(splitIntoUserInstalledAndSystemApps()).when(privateProfileManager)
+                .splitIntoUserInstalledAndSystemApps(any());
         doReturn(0).when(privateProfileManager).addPrivateSpaceHeader(any());
         doNothing().when(privateProfileManager).addPrivateSpaceInstallAppButton(any());
         doReturn(0).when(privateProfileManager).addSystemAppsDivider(any());
@@ -465,4 +458,10 @@
         }
         return appInfos.toArray(AppInfo[]::new);
     }
+
+    private Predicate<AppInfo> splitIntoUserInstalledAndSystemApps() {
+        return iteminfo -> iteminfo.componentName == null
+                || !iteminfo.componentName.getPackageName()
+                .equals(CAMERA_PACKAGE_NAME);
+    }
 }
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
index 342eedf..4a67600 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3Test.java
@@ -21,8 +21,9 @@
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.Launcher;
+import com.android.launcher3.util.rule.ScreenRecordRule.KeepRecordOnSuccess;
+import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -31,6 +32,7 @@
 @RunWith(AndroidJUnit4.class)
 public class TaplTestsLauncher3Test extends AbstractLauncherUiTest<Launcher> {
 
+    @KeepRecordOnSuccess
     @ScreenRecord // b/322823478
     @Test
     public void testDevicePressMenu() throws Exception {
diff --git a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
index e08d37c..70a3dd0 100644
--- a/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
+++ b/tests/src/com/android/launcher3/ui/TaplWorkProfileTest.java
@@ -45,6 +45,7 @@
 import com.android.launcher3.allapps.WorkProfileManager;
 import com.android.launcher3.tapl.LauncherInstrumentation;
 import com.android.launcher3.util.TestUtil;
+import com.android.launcher3.util.rule.ScreenRecordRule.KeepRecordOnSuccess;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.TestStabilityRule.Stability;
 
@@ -196,6 +197,7 @@
 
     }
 
+    @KeepRecordOnSuccess
     @ScreenRecord // b/322823478
     @Test
     public void testEdu() {
diff --git a/tests/src/com/android/launcher3/util/DisplayControllerTest.kt b/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
index 706ab27..2e57ad5 100644
--- a/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
+++ b/tests/src/com/android/launcher3/util/DisplayControllerTest.kt
@@ -45,6 +45,7 @@
 import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
@@ -123,6 +124,7 @@
         whenever(displayManager.getDisplay(any())).thenReturn(display)
 
         // Mock resources
+        doReturn(context).whenever(context).applicationContext
         whenever(resources.configuration).thenReturn(configuration)
         whenever(context.resources).thenReturn(resources)
 
diff --git a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
index 7a5cf2c..ff96afb 100644
--- a/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ScreenRecordRule.java
@@ -49,6 +49,9 @@
             return base;
         }
 
+        final Boolean keepRecordOnSuccess = description.getAnnotation(KeepRecordOnSuccess.class)
+                != null;
+
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
@@ -70,7 +73,7 @@
                     device.executeShellCommand("kill -INT " + screenRecordPid);
                     Log.e(TAG, "Screenrecord captured at: " + outputFile);
                     output.close();
-                    if (success) {
+                    if (success && !keepRecordOnSuccess) {
                         automation.executeShellCommand("rm " + outputFile);
                     }
                 }
@@ -85,4 +88,14 @@
     @Target(ElementType.METHOD)
     public @interface ScreenRecord {
     }
+
+
+    /**
+     * Interface to indicate that we should keep the screen record even if the test succeeds, use
+     * sparingly since it uses a lof of memory.
+     */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface KeepRecordOnSuccess {
+    }
 }