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 30bb863..2bb6b07 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -65,6 +65,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;
@@ -108,7 +109,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 =
@@ -199,6 +200,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/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/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index ec57115..13038b1 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -93,6 +93,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;
@@ -118,6 +119,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;
@@ -126,8 +129,6 @@
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
-import kotlin.Unit;
-
 /**
  * A task in the Recents view.
  */
@@ -922,8 +923,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
@@ -950,12 +951,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) {
@@ -1857,7 +1854,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;
@@ -1892,6 +1889,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/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/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/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 0e4b48e..9a4417d 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/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/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..b12c4eb 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() {
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/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/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/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/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)