[Launcher Jank] Avoid SimpleBroadcastReceiver making binder calls on main thread

Fix: 348649441
Flag: NONE - jank fix
Test: manual - presubmit
Change-Id: Ie97713f3f0b1f117662d49c6d6a65651c593d424
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 8b5ed7c..6af5a30 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -205,6 +205,7 @@
         mActive = true;
     }
 
+    @WorkerThread
     @Override
     public void workspaceLoadComplete() {
         super.workspaceLoadComplete();
@@ -323,6 +324,7 @@
         }
     }
 
+    @WorkerThread
     @Override
     public void destroy() {
         super.destroy();
diff --git a/quickstep/src/com/android/launcher3/model/WellbeingModel.java b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
index a7c9652..28bc01c 100644
--- a/quickstep/src/com/android/launcher3/model/WellbeingModel.java
+++ b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
@@ -111,6 +111,7 @@
         mWorkerHandler.post(this::initializeInBackground);
     }
 
+    @WorkerThread
     private void initializeInBackground() {
         if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
             mContext.registerReceiver(
@@ -134,8 +135,8 @@
     public void close() {
         if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
             mWorkerHandler.post(() -> {
-                mWellbeingAppChangeReceiver.unregisterReceiverSafely(mContext);
-                mAppAddRemoveReceiver.unregisterReceiverSafely(mContext);
+                mWellbeingAppChangeReceiver.unregisterReceiverSafelySync(mContext);
+                mAppAddRemoveReceiver.unregisterReceiverSafelySync(mContext);
                 mContext.getContentResolver().unregisterContentObserver(mContentObserver);
             });
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 2a58db2..051bdc8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -304,7 +304,7 @@
                 .register(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "registering component callbacks from constructor.");
         mContext.registerComponentCallbacks(mComponentCallbacks);
-        mShutdownReceiver.register(mContext, Intent.ACTION_SHUTDOWN);
+        mShutdownReceiver.registerAsync(mContext, Intent.ACTION_SHUTDOWN);
         UI_HELPER_EXECUTOR.execute(() -> {
             mSharedState.taskbarSystemActionPendingIntent = PendingIntent.getBroadcast(
                     mContext,
@@ -582,8 +582,7 @@
     public void destroy() {
         debugWhyTaskbarNotDestroyed("TaskbarManager#destroy()");
         removeActivityCallbacksAndListeners();
-        UI_HELPER_EXECUTOR.execute(
-                () -> mTaskbarBroadcastReceiver.unregisterReceiverSafely(mContext));
+        mTaskbarBroadcastReceiver.unregisterReceiverSafelyAsync(mContext);
         destroyExistingTaskbar();
         removeTaskbarRootViewFromWindow();
         if (mUserUnlocked) {
@@ -595,7 +594,7 @@
                 .unregister(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "unregistering component callbacks from destroy().");
         mContext.unregisterComponentCallbacks(mComponentCallbacks);
-        mContext.unregisterReceiver(mShutdownReceiver);
+        mShutdownReceiver.unregisterReceiverSafelyAsync(mContext);
     }
 
     public @Nullable TaskbarActivityContext getCurrentActivityContext() {
diff --git a/quickstep/src/com/android/quickstep/OverviewComponentObserver.java b/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
index a71e314..9c64576 100644
--- a/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
+++ b/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
@@ -36,6 +36,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.UiThread;
 
 import com.android.launcher3.R;
 import com.android.launcher3.util.SimpleBroadcastReceiver;
@@ -101,7 +102,7 @@
             mConfigChangesMap.append(fallbackComponent.hashCode(), fallbackInfo.configChanges);
         } catch (PackageManager.NameNotFoundException ignored) { /* Impossible */ }
 
-        mUserPreferenceChangeReceiver.register(mContext, ACTION_PREFERRED_ACTIVITY_CHANGED);
+        mUserPreferenceChangeReceiver.registerAsync(mContext, ACTION_PREFERRED_ACTIVITY_CHANGED);
         updateOverviewTargets();
     }
 
@@ -114,6 +115,8 @@
         mOverviewChangeListener = overviewChangeListener;
     }
 
+    /** Called on {@link TouchInteractionService#onSystemUiFlagsChanged} */
+    @UiThread
     public void onSystemUiStateChanged() {
         if (mDeviceState.isHomeDisabled() != mIsHomeDisabled) {
             updateOverviewTargets();
@@ -128,6 +131,7 @@
      * Update overview intent and {@link BaseActivityInterface} based off the current launcher home
      * component.
      */
+    @UiThread
     private void updateOverviewTargets() {
         ComponentName defaultHome = PackageManagerWrapper.getInstance()
                 .getHomeActivities(new ArrayList<>());
@@ -187,8 +191,9 @@
                 unregisterOtherHomeAppUpdateReceiver();
 
                 mUpdateRegisteredPackage = defaultHome.getPackageName();
-                mOtherHomeAppUpdateReceiver.registerPkgActions(mContext, mUpdateRegisteredPackage,
-                        ACTION_PACKAGE_ADDED, ACTION_PACKAGE_CHANGED, ACTION_PACKAGE_REMOVED);
+                mOtherHomeAppUpdateReceiver.registerPkgActionsAsync(
+                        mContext, mUpdateRegisteredPackage, ACTION_PACKAGE_ADDED,
+                        ACTION_PACKAGE_CHANGED, ACTION_PACKAGE_REMOVED);
             }
         }
         mOverviewChangeListener.accept(mIsHomeAndOverviewSame);
@@ -198,13 +203,13 @@
      * Clean up any registered receivers.
      */
     public void onDestroy() {
-        mContext.unregisterReceiver(mUserPreferenceChangeReceiver);
+        mUserPreferenceChangeReceiver.unregisterReceiverSafelyAsync(mContext);
         unregisterOtherHomeAppUpdateReceiver();
     }
 
     private void unregisterOtherHomeAppUpdateReceiver() {
         if (mUpdateRegisteredPackage != null) {
-            mContext.unregisterReceiver(mOtherHomeAppUpdateReceiver);
+            mOtherHomeAppUpdateReceiver.unregisterReceiverSafelyAsync(mContext);
             mUpdateRegisteredPackage = null;
         }
     }
diff --git a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
index cda87c0..c26fc0c5 100644
--- a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
+++ b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
@@ -18,8 +18,6 @@
 import static android.content.Intent.ACTION_TIMEZONE_CHANGED;
 import static android.content.Intent.ACTION_TIME_CHANGED;
 
-import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
-
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
@@ -64,9 +62,7 @@
     private AsyncClockEventDelegate(Context context) {
         super(context);
         mContext = context;
-
-        UI_HELPER_EXECUTOR.execute(() ->
-                mReceiver.register(mContext, ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED));
+        mReceiver.registerAsync(mContext, ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED);
     }
 
     @Override
@@ -127,6 +123,6 @@
     public void close() {
         mDestroyed = true;
         SettingsCache.INSTANCE.get(mContext).unregister(mFormatUri, this);
-        UI_HELPER_EXECUTOR.execute(() -> mReceiver.unregisterReceiverSafely(mContext));
+        mReceiver.unregisterReceiverSafelyAsync(mContext);
     }
 }
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 3b8ff62..239967d 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -116,13 +116,13 @@
 
         SimpleBroadcastReceiver modelChangeReceiver =
                 new SimpleBroadcastReceiver(mModel::onBroadcastIntent);
-        modelChangeReceiver.register(mContext, Intent.ACTION_LOCALE_CHANGED,
+        modelChangeReceiver.registerAsync(mContext, Intent.ACTION_LOCALE_CHANGED,
                 ACTION_DEVICE_POLICY_RESOURCE_UPDATED);
         if (BuildConfig.IS_STUDIO_BUILD) {
             mContext.registerReceiver(modelChangeReceiver, new IntentFilter(ACTION_FORCE_ROLOAD),
                     RECEIVER_EXPORTED);
         }
-        mOnTerminateCallback.add(() -> mContext.unregisterReceiver(modelChangeReceiver));
+        mOnTerminateCallback.add(() -> modelChangeReceiver.unregisterReceiverSafelyAsync(mContext));
 
         SafeCloseable userChangeListener = UserCache.INSTANCE.get(mContext)
                 .addUserEventListener(mModel::onUserEvent);
diff --git a/src/com/android/launcher3/pm/UserCache.java b/src/com/android/launcher3/pm/UserCache.java
index ed25186..cf03462 100644
--- a/src/com/android/launcher3/pm/UserCache.java
+++ b/src/com/android/launcher3/pm/UserCache.java
@@ -93,12 +93,12 @@
 
     @Override
     public void close() {
-        MODEL_EXECUTOR.execute(() -> mUserChangeReceiver.unregisterReceiverSafely(mContext));
+        MODEL_EXECUTOR.execute(() -> mUserChangeReceiver.unregisterReceiverSafelySync(mContext));
     }
 
     @WorkerThread
     private void initAsync() {
-        mUserChangeReceiver.register(mContext,
+        mUserChangeReceiver.registerSync(mContext,
                 Intent.ACTION_MANAGED_PROFILE_AVAILABLE,
                 Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE,
                 Intent.ACTION_MANAGED_PROFILE_REMOVED,
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 16fabe2..b390cb8 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -132,11 +132,11 @@
             mWindowContext.registerComponentCallbacks(this);
         } else {
             mWindowContext = null;
-            mReceiver.register(mContext, ACTION_CONFIGURATION_CHANGED);
+            mReceiver.registerAsync(mContext, ACTION_CONFIGURATION_CHANGED);
         }
 
         // Initialize navigation mode change listener
-        mReceiver.registerPkgActions(mContext, TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED);
+        mReceiver.registerPkgActionsAsync(mContext, TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED);
 
         WindowManagerProxy wmProxy = WindowManagerProxy.INSTANCE.get(context);
         Context displayInfoContext = getDisplayInfoContext(display);
@@ -218,6 +218,7 @@
         } else {
             // TODO: unregister broadcast receiver
         }
+        mReceiver.unregisterReceiverSafelyAsync(mContext);
     }
 
     /**
diff --git a/src/com/android/launcher3/util/LockedUserState.kt b/src/com/android/launcher3/util/LockedUserState.kt
index 94f9e4f..2737249 100644
--- a/src/com/android/launcher3/util/LockedUserState.kt
+++ b/src/com/android/launcher3/util/LockedUserState.kt
@@ -25,6 +25,7 @@
     val isUserUnlockedAtLauncherStartup: Boolean
     var isUserUnlocked: Boolean
         private set
+
     private val mUserUnlockedActions: RunnableList = RunnableList()
 
     @VisibleForTesting
@@ -50,22 +51,18 @@
         if (isUserUnlocked) {
             notifyUserUnlocked()
         } else {
-            mUserUnlockedReceiver.register(mContext, Intent.ACTION_USER_UNLOCKED)
+            mUserUnlockedReceiver.registerAsync(mContext, Intent.ACTION_USER_UNLOCKED)
         }
     }
 
     private fun notifyUserUnlocked() {
         mUserUnlockedActions.executeAllAndDestroy()
-        Executors.THREAD_POOL_EXECUTOR.execute {
-            mUserUnlockedReceiver.unregisterReceiverSafely(mContext)
-        }
+        mUserUnlockedReceiver.unregisterReceiverSafelyAsync(mContext)
     }
 
     /** Stops the receiver from listening for ACTION_USER_UNLOCK broadcasts. */
     override fun close() {
-        Executors.THREAD_POOL_EXECUTOR.execute {
-            mUserUnlockedReceiver.unregisterReceiverSafely(mContext)
-        }
+        mUserUnlockedReceiver.unregisterReceiverSafelyAsync(mContext)
     }
 
     /**
diff --git a/src/com/android/launcher3/util/ScreenOnTracker.java b/src/com/android/launcher3/util/ScreenOnTracker.java
index e16e477..c1d192c 100644
--- a/src/com/android/launcher3/util/ScreenOnTracker.java
+++ b/src/com/android/launcher3/util/ScreenOnTracker.java
@@ -42,12 +42,12 @@
         // 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);
+        mReceiver.registerAsync(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
     }
 
     @Override
     public void close() {
-        mReceiver.unregisterReceiverSafely(mContext);
+        mReceiver.unregisterReceiverSafelyAsync(mContext);
     }
 
     private void onReceive(Intent intent) {
diff --git a/src/com/android/launcher3/util/SimpleBroadcastReceiver.java b/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
index 064bcd0..5f39cce 100644
--- a/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
+++ b/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
@@ -15,14 +15,21 @@
  */
 package com.android.launcher3.util;
 
+import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
+
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.os.Looper;
 import android.os.PatternMatcher;
 import android.text.TextUtils;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.UiThread;
+import androidx.annotation.WorkerThread;
+
+import com.android.launcher3.BuildConfig;
 
 import java.util.function.Consumer;
 
@@ -39,21 +46,63 @@
         mIntentConsumer.accept(intent);
     }
 
-    /**
-     * Helper method to register multiple actions
-     */
-    public void register(Context context, String... actions) {
+    /** Helper method to register multiple actions. Caller should be on main thread. */
+    @UiThread
+    public void registerAsync(Context context, String... actions) {
+        assertOnMainThread();
+        UI_HELPER_EXECUTOR.execute(() -> registerSync(context, actions));
+    }
+
+    /** Helper method to register multiple actions. Caller should be on main thread. */
+    @WorkerThread
+    public void registerSync(Context context, String... actions) {
+        assertOnBgThread();
         context.registerReceiver(this, getFilter(actions));
     }
 
     /**
-     * Helper method to register multiple actions associated with a paction
+     * Helper method to register multiple actions associated with a action. Caller should be from
+     * main thread.
      */
-    public void registerPkgActions(Context context, @Nullable String pkg, String... actions) {
+    @UiThread
+    public void registerPkgActionsAsync(Context context, @Nullable String pkg, String... actions) {
+        assertOnMainThread();
+        UI_HELPER_EXECUTOR.execute(() -> registerPkgActionsSync(context, pkg, actions));
+    }
+
+    /**
+     * Helper method to register multiple actions associated with a action. Caller should be from
+     * bg thread.
+     */
+    @WorkerThread
+    public void registerPkgActionsSync(Context context, @Nullable String pkg, String... actions) {
+        assertOnBgThread();
         context.registerReceiver(this, getPackageFilter(pkg, actions));
     }
 
     /**
+     * Unregisters the receiver ignoring any errors on bg thread. Caller should be on main thread.
+     */
+    @UiThread
+    public void unregisterReceiverSafelyAsync(Context context) {
+        assertOnMainThread();
+        UI_HELPER_EXECUTOR.execute(() -> unregisterReceiverSafelySync(context));
+    }
+
+    /**
+     * Unregisters the receiver ignoring any errors on bg thread. Caller should be on bg thread.
+     */
+    @WorkerThread
+    public void unregisterReceiverSafelySync(Context context) {
+        assertOnBgThread();
+        try {
+            context.unregisterReceiver(this);
+        } catch (IllegalArgumentException e) {
+            // It was probably never registered or already unregistered. Ignore.
+        }
+    }
+
+    /**
      * Creates an intent filter to listen for actions with a specific package in the data field.
      */
     public static IntentFilter getPackageFilter(String pkg, String... actions) {
@@ -73,14 +122,19 @@
         return filter;
     }
 
-    /**
-     * Unregisters the receiver ignoring any errors
-     */
-    public void unregisterReceiverSafely(Context context) {
-        try {
-            context.unregisterReceiver(this);
-        } catch (IllegalArgumentException e) {
-            // It was probably never registered or already unregistered. Ignore.
+    private static void assertOnBgThread() {
+        if (BuildConfig.IS_STUDIO_BUILD && isMainThread()) {
+            throw new IllegalStateException("Should not be called from main thread!");
         }
     }
+
+    private static void assertOnMainThread() {
+        if (BuildConfig.IS_STUDIO_BUILD && !isMainThread()) {
+            throw new IllegalStateException("Should not be called from bg thread!");
+        }
+    }
+
+    private static boolean isMainThread() {
+        return Thread.currentThread() == Looper.getMainLooper().getThread();
+    }
 }
diff --git a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
index b97b889..a2277a0 100644
--- a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
+++ b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
@@ -198,10 +198,11 @@
     public void setWindowToken(IBinder token) {
         mWindowToken = token;
         if (mWindowToken == null && mRegistered) {
-            mWallpaperChangeReceiver.unregisterReceiverSafely(mWorkspace.getContext());
+            mWallpaperChangeReceiver.unregisterReceiverSafelyAsync(mWorkspace.getContext());
             mRegistered = false;
         } else if (mWindowToken != null && !mRegistered) {
-            mWallpaperChangeReceiver.register(mWorkspace.getContext(), ACTION_WALLPAPER_CHANGED);
+            mWallpaperChangeReceiver.registerAsync(
+                    mWorkspace.getContext(), ACTION_WALLPAPER_CHANGED);
             onWallpaperChanged();
             mRegistered = true;
         }
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index a6f4441..6e01f9e 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -239,7 +239,8 @@
         final CountDownLatch count = new CountDownLatch(2);
         final SimpleBroadcastReceiver broadcastReceiver =
                 new SimpleBroadcastReceiver(i -> count.countDown());
-        broadcastReceiver.registerPkgActions(mTargetContext, pkg,
+        // We OK to make binder calls on main thread in test.
+        broadcastReceiver.registerPkgActionsSync(mTargetContext, pkg,
                 Intent.ACTION_PACKAGE_RESTARTED, Intent.ACTION_PACKAGE_DATA_CLEARED);
 
         mDevice.executeShellCommand("pm clear " + pkg);