diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java
index 71aa2e8..381ecf1 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/touchcontrollers/NoButtonNavbarToOverviewTouchController.java
@@ -166,8 +166,8 @@
         float velocityDp = dpiFromPx(velocity);
         boolean isFling = Math.abs(velocityDp) > 1;
         LauncherStateManager stateManager = mLauncher.getStateManager();
-        if (isFling) {
-            // When flinging, go back to home instead of overview.
+        boolean goToHomeInsteadOfOverview = isFling;
+        if (goToHomeInsteadOfOverview) {
             if (velocity > 0) {
                 stateManager.goToState(NORMAL, true,
                         () -> onSwipeInteractionCompleted(NORMAL, Touch.FLING));
@@ -187,20 +187,21 @@
                         () -> onSwipeInteractionCompleted(NORMAL, Touch.SWIPE)));
                 anim.start();
             }
-        } else {
-            if (mReachedOverview) {
-                float distanceDp = dpiFromPx(Math.max(
-                        Math.abs(mRecentsView.getTranslationX()),
-                        Math.abs(mRecentsView.getTranslationY())));
-                long duration = (long) Math.max(TRANSLATION_ANIM_MIN_DURATION_MS,
-                        distanceDp / TRANSLATION_ANIM_VELOCITY_DP_PER_MS);
-                mRecentsView.animate()
-                        .translationX(0)
-                        .translationY(0)
-                        .setInterpolator(ACCEL_DEACCEL)
-                        .setDuration(duration)
-                        .withEndAction(this::maybeSwipeInteractionToOverviewComplete);
-            }
+        }
+        if (mReachedOverview) {
+            float distanceDp = dpiFromPx(Math.max(
+                    Math.abs(mRecentsView.getTranslationX()),
+                    Math.abs(mRecentsView.getTranslationY())));
+            long duration = (long) Math.max(TRANSLATION_ANIM_MIN_DURATION_MS,
+                    distanceDp / TRANSLATION_ANIM_VELOCITY_DP_PER_MS);
+            mRecentsView.animate()
+                    .translationX(0)
+                    .translationY(0)
+                    .setInterpolator(ACCEL_DEACCEL)
+                    .setDuration(duration)
+                    .withEndAction(goToHomeInsteadOfOverview
+                            ? null
+                            : this::maybeSwipeInteractionToOverviewComplete);
         }
     }
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
index bf01780..52b40a9 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherSwipeHandler.java
@@ -1077,7 +1077,8 @@
     }
 
     private void continueComputingRecentsScrollIfNecessary() {
-        if (!mGestureState.hasState(STATE_RECENTS_SCROLLING_FINISHED)) {
+        if (!mGestureState.hasState(STATE_RECENTS_SCROLLING_FINISHED)
+                && !mStateCallback.hasStates(STATE_HANDLER_INVALIDATED)) {
             computeRecentsScrollIfInvisible();
             mRecentsView.post(this::continueComputingRecentsScrollIfNecessary);
         }
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index 65763d4..af63a25 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -161,6 +161,8 @@
     @Override
     protected void setupViews() {
         super.setupViews();
+
+        SysUINavigationMode.INSTANCE.get(this).updateMode();
         mActionsView = findViewById(R.id.overview_actions_view);
         ((RecentsView) getOverviewPanel()).init(mActionsView);
 
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java
index cc3fd97..e5c9fc9 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/PortraitStatesTouchController.java
@@ -22,10 +22,12 @@
 import static com.android.launcher3.anim.Interpolators.ACCEL;
 import static com.android.launcher3.anim.Interpolators.DEACCEL;
 import static com.android.launcher3.anim.Interpolators.LINEAR;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_OVERVIEW_ACTIONS;
 import static com.android.launcher3.config.FeatureFlags.UNSTABLE_SPRINGS;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_ALL_APPS_FADE;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_OVERVIEW_FADE;
 import static com.android.launcher3.states.StateAnimationConfig.ANIM_VERTICAL_PROGRESS;
+import static com.android.quickstep.SysUINavigationMode.removeShelfFromOverview;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_OVERVIEW_DISABLED;
 
 import android.animation.TimeInterpolator;
@@ -132,7 +134,12 @@
             return TouchInteractionService.isConnected() ?
                     mLauncher.getStateManager().getLastState() : NORMAL;
         } else if (fromState == OVERVIEW) {
-            return isDragTowardPositive ? ALL_APPS : NORMAL;
+            LauncherState positiveDragTarget = ALL_APPS;
+            if (ENABLE_OVERVIEW_ACTIONS.get() && removeShelfFromOverview(mLauncher)) {
+                // Don't allow swiping up to all apps.
+                positiveDragTarget = OVERVIEW;
+            }
+            return isDragTowardPositive ? positiveDragTarget : NORMAL;
         } else if (fromState == NORMAL && isDragTowardPositive) {
             int stateFlags = SystemUiProxy.INSTANCE.get(mLauncher).getLastSystemUiStateFlags();
             return mAllowDragToOverview && TouchInteractionService.isConnected()
diff --git a/quickstep/src/com/android/quickstep/SysUINavigationMode.java b/quickstep/src/com/android/quickstep/SysUINavigationMode.java
index 6a10b37..c715c93 100644
--- a/quickstep/src/com/android/quickstep/SysUINavigationMode.java
+++ b/quickstep/src/com/android/quickstep/SysUINavigationMode.java
@@ -74,15 +74,20 @@
         mContext.registerReceiver(new BroadcastReceiver() {
             @Override
             public void onReceive(Context context, Intent intent) {
-                Mode oldMode = mMode;
-                initializeMode();
-                if (mMode != oldMode) {
-                    dispatchModeChange();
-                }
+                updateMode();
             }
         }, getPackageFilter("android", ACTION_OVERLAY_CHANGED));
     }
 
+    /** Updates navigation mode when needed. */
+    public void updateMode() {
+        Mode oldMode = mMode;
+        initializeMode();
+        if (mMode != oldMode) {
+            dispatchModeChange();
+        }
+    }
+
     private void initializeMode() {
         int modeInt = getSystemIntegerRes(mContext, NAV_BAR_INTERACTION_MODE_RES_NAME);
         for(Mode m : Mode.values()) {
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 0afe4a8..40265c4 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -203,7 +203,7 @@
                         + launcher.getNavigationModeMismatchError(),
                 () -> launcher.getNavigationModeMismatchError() == null,
                 60000 /* b/148422894 */, launcher);
-        AbstractLauncherUiTest.checkDetectedLeaks();
+        AbstractLauncherUiTest.checkDetectedLeaks(launcher);
         return true;
     }
 
diff --git a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
index b7340cf..bbbe21e 100644
--- a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -130,7 +130,7 @@
         }
         helper.close();
 
-        helper = new DatabaseHelper(mContext, DB_FILE) {
+        helper = new DatabaseHelper(mContext, DB_FILE, false) {
             @Override
             public void onOpen(SQLiteDatabase db) { }
         };
@@ -161,7 +161,7 @@
 
         DbDowngradeHelper.updateSchemaFile(mSchemaFile, LauncherProvider.SCHEMA_VERSION, mContext);
 
-        DatabaseHelper dbHelper = new DatabaseHelper(mContext, DB_FILE) {
+        DatabaseHelper dbHelper = new DatabaseHelper(mContext, DB_FILE, false) {
             @Override
             public void onOpen(SQLiteDatabase db) { }
         };
diff --git a/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 58174c7..ee73b82 100644
--- a/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/robolectric_tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -95,7 +95,7 @@
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(RuntimeEnvironment.application, null);
+            super(RuntimeEnvironment.application, null, false);
             mProfileId = profileId;
         }
 
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 0a1fad1..63b90ae 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -17,6 +17,7 @@
 package com.android.launcher3;
 
 import static com.android.launcher3.Utilities.getDevicePrefs;
+import static com.android.launcher3.Utilities.getPointString;
 import static com.android.launcher3.config.FeatureFlags.APPLY_CONFIG_AT_RUNTIME;
 import static com.android.launcher3.settings.SettingsActivity.GRID_OPTIONS_PREFERENCE_KEY;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
@@ -69,6 +70,9 @@
     public static final MainThreadInitializedObject<InvariantDeviceProfile> INSTANCE =
             new MainThreadInitializedObject<>(InvariantDeviceProfile::new);
 
+    public static final String KEY_MIGRATION_SRC_WORKSPACE_SIZE = "migration_src_workspace_size";
+    public static final String KEY_MIGRATION_SRC_HOTSEAT_COUNT = "migration_src_hotseat_count";
+
     private static final String KEY_IDP_GRID_NAME = "idp_grid_name";
 
     private static final float ICON_SIZE_DEFINED_IN_APP_DP = 48;
@@ -165,6 +169,10 @@
         if (!newGridName.equals(gridName)) {
             Utilities.getPrefs(context).edit().putString(KEY_IDP_GRID_NAME, newGridName).apply();
         }
+        Utilities.getPrefs(context).edit()
+                .putInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, numHotseatIcons)
+                .putString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, getPointString(numColumns, numRows))
+                .apply();
 
         mConfigMonitor = new ConfigMonitor(context,
                 APPLY_CONFIG_AT_RUNTIME.get() ? this::onConfigChanged : this::killProcess);
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 8d20bd6..308d84f 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -146,7 +146,8 @@
      */
     protected synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
-            mOpenHelper = DatabaseHelper.createDatabaseHelper(getContext());
+            mOpenHelper = DatabaseHelper.createDatabaseHelper(
+                    getContext(), false /* forMigration */);
 
             if (RestoreDbTask.isPending(getContext())) {
                 if (!RestoreDbTask.performRestore(getContext(), mOpenHelper,
@@ -430,7 +431,8 @@
                                     InvariantDeviceProfile.INSTANCE.get(getContext()).dbFile,
                                     Favorites.TMP_TABLE,
                                     () -> mOpenHelper,
-                                    () -> DatabaseHelper.createDatabaseHelper(getContext())));
+                                    () -> DatabaseHelper.createDatabaseHelper(
+                                            getContext(), true /* forMigration */)));
                     return result;
                 }
             }
@@ -441,7 +443,8 @@
                             prepForMigration(
                                     arg /* dbFile */,
                                     Favorites.PREVIEW_TABLE_NAME,
-                                    () -> DatabaseHelper.createDatabaseHelper(getContext(), arg),
+                                    () -> DatabaseHelper.createDatabaseHelper(
+                                            getContext(), arg, true /* forMigration */),
                                     () -> mOpenHelper));
                     return result;
                 }
@@ -609,20 +612,22 @@
     public static class DatabaseHelper extends NoLocaleSQLiteHelper implements
             LayoutParserCallback {
         private final Context mContext;
+        private final boolean mForMigration;
         private int mMaxItemId = -1;
         private int mMaxScreenId = -1;
         private boolean mBackupTableExists;
 
-        static DatabaseHelper createDatabaseHelper(Context context) {
-            return createDatabaseHelper(context, null);
+        static DatabaseHelper createDatabaseHelper(Context context, boolean forMigration) {
+            return createDatabaseHelper(context, null, forMigration);
         }
 
-        static DatabaseHelper createDatabaseHelper(Context context, String dbName) {
+        static DatabaseHelper createDatabaseHelper(Context context, String dbName,
+                boolean forMigration) {
             if (dbName == null) {
                 dbName = MULTI_DB_GRID_MIRATION_ALGO.get() ? InvariantDeviceProfile.INSTANCE.get(
                         context).dbFile : LauncherFiles.LAUNCHER_DB;
             }
-            DatabaseHelper databaseHelper = new DatabaseHelper(context, dbName);
+            DatabaseHelper databaseHelper = new DatabaseHelper(context, dbName, forMigration);
             // Table creation sometimes fails silently, which leads to a crash loop.
             // This way, we will try to create a table every time after crash, so the device
             // would eventually be able to recover.
@@ -643,9 +648,10 @@
         /**
          * Constructor used in tests and for restore.
          */
-        public DatabaseHelper(Context context, String dbName) {
+        public DatabaseHelper(Context context, String dbName, boolean forMigration) {
             super(context, dbName, SCHEMA_VERSION);
             mContext = context;
+            mForMigration = forMigration;
         }
 
         protected void initIds() {
@@ -670,7 +676,9 @@
 
             // Fresh and clean launcher DB.
             mMaxItemId = initializeMaxItemId(db);
-            onEmptyDbCreated();
+            if (!mForMigration) {
+                onEmptyDbCreated();
+            }
         }
 
         protected void onAddOrDeleteOp(SQLiteDatabase db) {
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTask.java b/src/com/android/launcher3/model/GridSizeMigrationTask.java
index b27e4ea..e8a52bd 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTask.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTask.java
@@ -1,5 +1,7 @@
 package com.android.launcher3.model;
 
+import static com.android.launcher3.InvariantDeviceProfile.KEY_MIGRATION_SRC_HOTSEAT_COUNT;
+import static com.android.launcher3.InvariantDeviceProfile.KEY_MIGRATION_SRC_WORKSPACE_SIZE;
 import static com.android.launcher3.LauncherSettings.Settings.EXTRA_VALUE;
 import static com.android.launcher3.Utilities.getPointString;
 import static com.android.launcher3.Utilities.parsePoint;
@@ -53,9 +55,6 @@
     private static final String TAG = "GridSizeMigrationTask";
     private static final boolean DEBUG = true;
 
-    private static final String KEY_MIGRATION_SRC_WORKSPACE_SIZE = "migration_src_workspace_size";
-    private static final String KEY_MIGRATION_SRC_HOTSEAT_COUNT = "migration_src_hotseat_count";
-
     // These are carefully selected weights for various item types (Math.random?), to allow for
     // the least absurd migration experience.
     private static final float WT_SHORTCUT = 1;
@@ -894,8 +893,7 @@
         String gridSizeString = getPointString(idp.numColumns, idp.numRows);
 
         return !gridSizeString.equals(prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, ""))
-                || idp.numHotseatIcons != prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT,
-                idp.numHotseatIcons);
+                || idp.numHotseatIcons != prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, -1);
     }
 
     /** See {@link #migrateGridIfNeeded(Context, InvariantDeviceProfile)} */
diff --git a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
index 1c44fc3..4a28218 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationTaskV2.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.model;
 
+import static com.android.launcher3.InvariantDeviceProfile.KEY_MIGRATION_SRC_HOTSEAT_COUNT;
+import static com.android.launcher3.InvariantDeviceProfile.KEY_MIGRATION_SRC_WORKSPACE_SIZE;
 import static com.android.launcher3.Utilities.getPointString;
 import static com.android.launcher3.provider.LauncherDbUtils.dropTable;
 
@@ -63,9 +65,6 @@
  */
 public class GridSizeMigrationTaskV2 {
 
-    public static final String KEY_MIGRATION_SRC_WORKSPACE_SIZE = "migration_src_workspace_size";
-    public static final String KEY_MIGRATION_SRC_HOTSEAT_COUNT = "migration_src_hotseat_count";
-
     private static final String TAG = "GridSizeMigrationTaskV2";
     private static final boolean DEBUG = true;
 
@@ -110,8 +109,7 @@
         String gridSizeString = getPointString(idp.numColumns, idp.numRows);
 
         return !gridSizeString.equals(prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, ""))
-                || idp.numHotseatIcons != prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT,
-                idp.numHotseatIcons);
+                || idp.numHotseatIcons != prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, -1);
     }
 
     /** See {@link #migrateGridIfNeeded(Context, InvariantDeviceProfile)} */
@@ -148,14 +146,6 @@
 
         SharedPreferences prefs = Utilities.getPrefs(context);
         String gridSizeString = getPointString(idp.numColumns, idp.numRows);
-
-        if (gridSizeString.equals(prefs.getString(KEY_MIGRATION_SRC_WORKSPACE_SIZE, ""))
-                && idp.numHotseatIcons == prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT,
-                idp.numHotseatIcons)) {
-            // Skip if workspace and hotseat sizes have not changed.
-            return true;
-        }
-
         HashSet<String> validPackages = getValidPackages(context);
         int srcHotseatCount = prefs.getInt(KEY_MIGRATION_SRC_HOTSEAT_COUNT, idp.numHotseatIcons);
 
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 0ba5a36..9c8e278 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -98,9 +98,10 @@
     public static final long DEFAULT_UI_TIMEOUT = 10000;
     private static final String TAG = "AbstractLauncherUiTest";
 
-    private static String sDetectedActivityLeak;
+    private static String sStrictmodeDetectedActivityLeak;
     private static boolean sActivityLeakReported;
     private static final String SYSTEMUI_PACKAGE = "com.android.systemui";
+    private static final ActivityLeakTracker ACTIVITY_LEAK_TRACKER = new ActivityLeakTracker();
 
     protected LooperExecutor mMainThreadExecutor = MAIN_EXECUTOR;
     protected final UiDevice mDevice = UiDevice.getInstance(getInstrumentation());
@@ -113,30 +114,51 @@
         if (TestHelpers.isInLauncherProcess()) {
             StrictMode.VmPolicy.Builder builder =
                     new StrictMode.VmPolicy.Builder()
-                            .detectActivityLeaks()
+// b/154772063
+//                            .detectActivityLeaks()
                             .penaltyLog()
                             .penaltyListener(Runnable::run, violation -> {
-                                // Runs in the main thread. We can't dumpheap in the main thread,
-                                // so let's just mark the fact that the leak has happened.
-                                if (sDetectedActivityLeak == null) {
-                                    sDetectedActivityLeak = violation.toString();
-                                    try {
-                                        Debug.dumpHprofData(
-                                                getInstrumentation().getTargetContext()
-                                                        .getFilesDir().getPath()
-                                                        + "/ActivityLeakHeapDump.hprof");
-                                    } catch (Throwable e) {
-                                        Log.e(TAG, "dumpHprofData failed", e);
-                                    }
+                                if (sStrictmodeDetectedActivityLeak == null) {
+                                    sStrictmodeDetectedActivityLeak = violation.toString() + ", "
+                                            + dumpHprofData() + ".";
                                 }
                             });
             StrictMode.setVmPolicy(builder.build());
         }
     }
 
-    public static void checkDetectedLeaks() {
-        if (sDetectedActivityLeak != null && !sActivityLeakReported) {
+    public static void checkDetectedLeaks(LauncherInstrumentation launcher) {
+        if (sActivityLeakReported) return;
+
+        if (sStrictmodeDetectedActivityLeak != null) {
+            // Report from the test thread strictmode violations detected in the main thread.
             sActivityLeakReported = true;
+            Assert.fail(sStrictmodeDetectedActivityLeak);
+        }
+
+        // Check whether activity leak detector has found leaked activities.
+        Wait.atMost(AbstractLauncherUiTest::getActivityLeakErrorMessage,
+                () -> {
+                    launcher.getTotalPssKb();  // Triggers GC
+                    return MAIN_EXECUTOR.submit(
+                            () -> ACTIVITY_LEAK_TRACKER.noLeakedActivities()).get();
+                }, DEFAULT_UI_TIMEOUT, launcher);
+    }
+
+    private static String getActivityLeakErrorMessage() {
+        sActivityLeakReported = true;
+        return "Activity leak detector has found leaked activities, " + dumpHprofData() + ".";
+    }
+
+    private static String dumpHprofData() {
+        try {
+            final String fileName = getInstrumentation().getTargetContext().getFilesDir().getPath()
+                    + "/ActivityLeakHeapDump.hprof";
+            Debug.dumpHprofData(fileName);
+            return "memory dump filename: " + fileName;
+        } catch (Throwable e) {
+            Log.e(TAG, "dumpHprofData failed", e);
+            return "failed to save memory dump";
         }
     }
 
@@ -263,7 +285,7 @@
         if (mLauncherPid != 0) {
             assertEquals("Launcher crashed, pid mismatch:", mLauncherPid, mLauncher.getPid());
         }
-        checkDetectedLeaks();
+        checkDetectedLeaks(mLauncher);
     }
 
     protected void clearLauncherData() throws IOException, InterruptedException {
diff --git a/tests/src/com/android/launcher3/ui/ActivityLeakTracker.java b/tests/src/com/android/launcher3/ui/ActivityLeakTracker.java
new file mode 100644
index 0000000..e9258e9
--- /dev/null
+++ b/tests/src/com/android/launcher3/ui/ActivityLeakTracker.java
@@ -0,0 +1,85 @@
+/*
+ * 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.ui;
+
+import android.app.Activity;
+import android.app.Application;
+import android.os.Bundle;
+
+import androidx.test.InstrumentationRegistry;
+
+import com.android.launcher3.tapl.TestHelpers;
+
+import java.util.WeakHashMap;
+
+class ActivityLeakTracker implements Application.ActivityLifecycleCallbacks {
+    private final WeakHashMap<Activity, Boolean> mActivities = new WeakHashMap<>();
+
+    ActivityLeakTracker() {
+        if (!TestHelpers.isInLauncherProcess()) return;
+        final Application app =
+                (Application) InstrumentationRegistry.getTargetContext().getApplicationContext();
+        app.registerActivityLifecycleCallbacks(this);
+    }
+
+    @Override
+    public void onActivityCreated(Activity activity, Bundle bundle) {
+        mActivities.put(activity, true);
+    }
+
+    @Override
+    public void onActivityStarted(Activity activity) {
+    }
+
+    @Override
+    public void onActivityResumed(Activity activity) {
+    }
+
+    @Override
+    public void onActivityPaused(Activity activity) {
+    }
+
+    @Override
+    public void onActivityStopped(Activity activity) {
+    }
+
+    @Override
+    public void onActivitySaveInstanceState(Activity activity, Bundle bundle) {
+    }
+
+    @Override
+    public void onActivityDestroyed(Activity activity) {
+    }
+
+    public boolean noLeakedActivities() {
+        int liveActivities = 0;
+        int destroyedActivities = 0;
+
+        for (Activity activity : mActivities.keySet()) {
+            if (activity.isDestroyed()) {
+                ++destroyedActivities;
+            } else {
+                ++liveActivities;
+            }
+        }
+
+        if (liveActivities > 2)  return false;
+
+        // It's OK to have 1 leaked activity if no active activities exist.
+        return liveActivities == 0 ? destroyedActivities <= 1 : destroyedActivities == 0;
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
index 38f50c1..266f0ae 100644
--- a/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
+++ b/tests/src/com/android/launcher3/ui/PortraitLandscapeRunner.java
@@ -56,7 +56,7 @@
             private void evaluateInPortrait() throws Throwable {
                 mTest.mDevice.setOrientationNatural();
                 mTest.mLauncher.setExpectedRotation(Surface.ROTATION_0);
-                AbstractLauncherUiTest.checkDetectedLeaks();
+                AbstractLauncherUiTest.checkDetectedLeaks(mTest.mLauncher);
                 base.evaluate();
                 mTest.getDevice().pressHome();
             }
@@ -64,7 +64,7 @@
             private void evaluateInLandscape() throws Throwable {
                 mTest.mDevice.setOrientationLeft();
                 mTest.mLauncher.setExpectedRotation(Surface.ROTATION_90);
-                AbstractLauncherUiTest.checkDetectedLeaks();
+                AbstractLauncherUiTest.checkDetectedLeaks(mTest.mLauncher);
                 base.evaluate();
                 mTest.getDevice().pressHome();
             }
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 57000a0..34e425d 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -64,7 +64,7 @@
         test.waitForResumed("Launcher internal state is still Background");
         // Check that we switched to home.
         test.mLauncher.getWorkspace();
-        AbstractLauncherUiTest.checkDetectedLeaks();
+        AbstractLauncherUiTest.checkDetectedLeaks(test.mLauncher);
     }
 
     // Please don't add negative test cases for methods that fail only after a long wait.