[Launcher Jank] Improve SimpleBroadcastReceiver.java

Bug: 348649441
Flag: NONE - jank fix
Test: manual - presubmit
Change-Id: I17bd7e4d7f0640522476e94edae691f983835f62
diff --git a/quickstep/src/com/android/launcher3/model/WellbeingModel.java b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
index 28bc01c..fb17f15 100644
--- a/quickstep/src/com/android/launcher3/model/WellbeingModel.java
+++ b/quickstep/src/com/android/launcher3/model/WellbeingModel.java
@@ -83,10 +83,8 @@
 
     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 SimpleBroadcastReceiver mWellbeingAppChangeReceiver;
+    private final SimpleBroadcastReceiver mAppAddRemoveReceiver;
 
     private final Object mModelLock = new Object();
     // Maps the action Id to the corresponding RemoteAction
@@ -101,6 +99,11 @@
         mWorkerHandler = new Handler(TextUtils.isEmpty(mWellbeingProviderPkg)
                 ? Executors.UI_HELPER_EXECUTOR.getLooper()
                 : Executors.getPackageExecutor(mWellbeingProviderPkg).getLooper());
+        mWellbeingAppChangeReceiver =
+                new SimpleBroadcastReceiver(mWorkerHandler, t -> restartObserver());
+        mAppAddRemoveReceiver =
+                new SimpleBroadcastReceiver(mWorkerHandler, this::onAppPackageChanged);
+
 
         mContentObserver = new ContentObserver(mWorkerHandler) {
             @Override
@@ -135,8 +138,8 @@
     public void close() {
         if (!TextUtils.isEmpty(mWellbeingProviderPkg)) {
             mWorkerHandler.post(() -> {
-                mWellbeingAppChangeReceiver.unregisterReceiverSafelySync(mContext);
-                mAppAddRemoveReceiver.unregisterReceiverSafelySync(mContext);
+                mWellbeingAppChangeReceiver.unregisterReceiverSafely(mContext);
+                mAppAddRemoveReceiver.unregisterReceiverSafely(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 b90e5fd..f411e79 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -39,7 +39,6 @@
 import android.content.ComponentCallbacks;
 import android.content.Context;
 import android.content.Intent;
-import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
 import android.content.res.Configuration;
 import android.hardware.display.DisplayManager;
@@ -120,7 +119,7 @@
     private final ComponentCallbacks mComponentCallbacks;
 
     private final SimpleBroadcastReceiver mShutdownReceiver =
-            new SimpleBroadcastReceiver(i -> destroyExistingTaskbar());
+            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, i -> destroyExistingTaskbar());
 
     // The source for this provider is set when Launcher is available
     // We use 'non-destroyable' version here so the original provider won't be destroyed
@@ -157,7 +156,7 @@
     private boolean mUserUnlocked = false;
 
     private final SimpleBroadcastReceiver mTaskbarBroadcastReceiver =
-            new SimpleBroadcastReceiver(this::showTaskbarFromBroadcast);
+            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::showTaskbarFromBroadcast);
 
     private final AllAppsActionManager mAllAppsActionManager;
 
@@ -306,17 +305,15 @@
                 .register(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "registering component callbacks from constructor.");
         mContext.registerComponentCallbacks(mComponentCallbacks);
-        mShutdownReceiver.registerAsync(mContext, Intent.ACTION_SHUTDOWN);
+        mShutdownReceiver.register(mContext, Intent.ACTION_SHUTDOWN);
         UI_HELPER_EXECUTOR.execute(() -> {
             mSharedState.taskbarSystemActionPendingIntent = PendingIntent.getBroadcast(
                     mContext,
                     SYSTEM_ACTION_ID_TASKBAR,
                     new Intent(ACTION_SHOW_TASKBAR).setPackage(mContext.getPackageName()),
                     PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
-            mContext.registerReceiver(
-                    mTaskbarBroadcastReceiver,
-                    new IntentFilter(ACTION_SHOW_TASKBAR),
-                    RECEIVER_NOT_EXPORTED);
+            mTaskbarBroadcastReceiver.register(
+                    mContext, RECEIVER_NOT_EXPORTED, ACTION_SHOW_TASKBAR);
         });
 
         debugWhyTaskbarNotDestroyed("TaskbarManager created");
@@ -623,7 +620,7 @@
     public void destroy() {
         debugWhyTaskbarNotDestroyed("TaskbarManager#destroy()");
         removeActivityCallbacksAndListeners();
-        mTaskbarBroadcastReceiver.unregisterReceiverSafelyAsync(mContext);
+        mTaskbarBroadcastReceiver.unregisterReceiverSafely(mContext);
         destroyExistingTaskbar();
         removeTaskbarRootViewFromWindow();
         if (mUserUnlocked) {
@@ -635,7 +632,7 @@
                 .unregister(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "unregistering component callbacks from destroy().");
         mContext.unregisterComponentCallbacks(mComponentCallbacks);
-        mShutdownReceiver.unregisterReceiverSafelyAsync(mContext);
+        mShutdownReceiver.unregisterReceiverSafely(mContext);
     }
 
     public @Nullable TaskbarActivityContext getCurrentActivityContext() {
diff --git a/quickstep/src/com/android/quickstep/OverviewComponentObserver.java b/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
index 9c64576..d82426f 100644
--- a/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
+++ b/quickstep/src/com/android/quickstep/OverviewComponentObserver.java
@@ -21,6 +21,7 @@
 import static android.content.Intent.ACTION_PACKAGE_REMOVED;
 
 import static com.android.launcher3.config.FeatureFlags.SEPARATE_RECENTS_ACTIVITY;
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.systemui.shared.system.PackageManagerWrapper.ACTION_PREFERRED_ACTIVITY_CHANGED;
 
 import android.content.ActivityNotFoundException;
@@ -55,10 +56,11 @@
 public final class OverviewComponentObserver {
     private static final String TAG = "OverviewComponentObserver";
 
+    // We register broadcast receivers on main thread to avoid missing updates.
     private final SimpleBroadcastReceiver mUserPreferenceChangeReceiver =
-            new SimpleBroadcastReceiver(this::updateOverviewTargets);
+            new SimpleBroadcastReceiver(MAIN_EXECUTOR, this::updateOverviewTargets);
     private final SimpleBroadcastReceiver mOtherHomeAppUpdateReceiver =
-            new SimpleBroadcastReceiver(this::updateOverviewTargets);
+            new SimpleBroadcastReceiver(MAIN_EXECUTOR, this::updateOverviewTargets);
 
     private final Context mContext;
     private final RecentsAnimationDeviceState mDeviceState;
@@ -102,7 +104,7 @@
             mConfigChangesMap.append(fallbackComponent.hashCode(), fallbackInfo.configChanges);
         } catch (PackageManager.NameNotFoundException ignored) { /* Impossible */ }
 
-        mUserPreferenceChangeReceiver.registerAsync(mContext, ACTION_PREFERRED_ACTIVITY_CHANGED);
+        mUserPreferenceChangeReceiver.register(mContext, ACTION_PREFERRED_ACTIVITY_CHANGED);
         updateOverviewTargets();
     }
 
@@ -191,7 +193,7 @@
                 unregisterOtherHomeAppUpdateReceiver();
 
                 mUpdateRegisteredPackage = defaultHome.getPackageName();
-                mOtherHomeAppUpdateReceiver.registerPkgActionsAsync(
+                mOtherHomeAppUpdateReceiver.registerPkgActions(
                         mContext, mUpdateRegisteredPackage, ACTION_PACKAGE_ADDED,
                         ACTION_PACKAGE_CHANGED, ACTION_PACKAGE_REMOVED);
             }
@@ -203,13 +205,13 @@
      * Clean up any registered receivers.
      */
     public void onDestroy() {
-        mUserPreferenceChangeReceiver.unregisterReceiverSafelyAsync(mContext);
+        mUserPreferenceChangeReceiver.unregisterReceiverSafely(mContext);
         unregisterOtherHomeAppUpdateReceiver();
     }
 
     private void unregisterOtherHomeAppUpdateReceiver() {
         if (mUpdateRegisteredPackage != null) {
-            mOtherHomeAppUpdateReceiver.unregisterReceiverSafelyAsync(mContext);
+            mOtherHomeAppUpdateReceiver.unregisterReceiverSafely(mContext);
             mUpdateRegisteredPackage = null;
         }
     }
diff --git a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
index c26fc0c5..38ae303 100644
--- a/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
+++ b/quickstep/src/com/android/quickstep/util/AsyncClockEventDelegate.java
@@ -18,6 +18,8 @@
 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;
@@ -50,7 +52,7 @@
 
     private final Context mContext;
     private final SimpleBroadcastReceiver mReceiver =
-            new SimpleBroadcastReceiver(this::onClockEventReceived);
+            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::onClockEventReceived);
 
     private final ArrayMap<BroadcastReceiver, Handler> mTimeEventReceivers = new ArrayMap<>();
     private final List<ContentObserver> mFormatObservers = new ArrayList<>();
@@ -62,7 +64,7 @@
     private AsyncClockEventDelegate(Context context) {
         super(context);
         mContext = context;
-        mReceiver.registerAsync(mContext, ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED);
+        mReceiver.register(mContext, ACTION_TIME_CHANGED, ACTION_TIMEZONE_CHANGED);
     }
 
     @Override
@@ -123,6 +125,6 @@
     public void close() {
         mDestroyed = true;
         SettingsCache.INSTANCE.get(mContext).unregister(mFormatUri, this);
-        mReceiver.unregisterReceiverSafelyAsync(mContext);
+        mReceiver.unregisterReceiverSafely(mContext);
     }
 }
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 239967d..85c8b57 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -24,6 +24,7 @@
 import static com.android.launcher3.LauncherPrefs.THEMED_ICONS;
 import static com.android.launcher3.model.LoaderTask.SMARTSPACE_ON_HOME_SCREEN;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
+import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.SettingsCache.NOTIFICATION_BADGING_URI;
 import static com.android.launcher3.util.SettingsCache.PRIVATE_SPACE_HIDE_WHEN_LOCKED_URI;
 
@@ -63,6 +64,9 @@
 import com.android.launcher3.util.TraceHelper;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
+import java.util.Locale;
+import java.util.Objects;
+
 public class LauncherAppState implements SafeCloseable {
 
     public static final String ACTION_FORCE_ROLOAD = "force-reload-launcher";
@@ -115,14 +119,25 @@
         }
 
         SimpleBroadcastReceiver modelChangeReceiver =
-                new SimpleBroadcastReceiver(mModel::onBroadcastIntent);
-        modelChangeReceiver.registerAsync(mContext, Intent.ACTION_LOCALE_CHANGED,
+                new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, mModel::onBroadcastIntent);
+        final Locale oldLocale = mContext.getResources().getConfiguration().locale;
+        modelChangeReceiver.register(
+                mContext,
+                () -> {
+                    // if local has changed before receiver is registered on bg thread,
+                    // mModel needs to reload.
+                    Locale newLocale = mContext.getResources().getConfiguration().locale;
+                    if (!Objects.equals(oldLocale, newLocale)) {
+                        mModel.forceReload();
+                    }
+                },
+                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(() -> modelChangeReceiver.unregisterReceiverSafelyAsync(mContext));
+        mOnTerminateCallback.add(() -> modelChangeReceiver.unregisterReceiverSafely(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 cf03462..7339111 100644
--- a/src/com/android/launcher3/pm/UserCache.java
+++ b/src/com/android/launcher3/pm/UserCache.java
@@ -75,7 +75,7 @@
 
     private final List<BiConsumer<UserHandle, String>> mUserEventListeners = new ArrayList<>();
     private final SimpleBroadcastReceiver mUserChangeReceiver =
-            new SimpleBroadcastReceiver(this::onUsersChanged);
+            new SimpleBroadcastReceiver(MODEL_EXECUTOR, this::onUsersChanged);
 
     private final Context mContext;
 
@@ -93,12 +93,12 @@
 
     @Override
     public void close() {
-        MODEL_EXECUTOR.execute(() -> mUserChangeReceiver.unregisterReceiverSafelySync(mContext));
+        MODEL_EXECUTOR.execute(() -> mUserChangeReceiver.unregisterReceiverSafely(mContext));
     }
 
     @WorkerThread
     private void initAsync() {
-        mUserChangeReceiver.registerSync(mContext,
+        mUserChangeReceiver.register(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 3dcc663..8556ae6 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -109,7 +109,10 @@
     private DisplayInfoChangeListener mPriorityListener;
     private final ArrayList<DisplayInfoChangeListener> mListeners = new ArrayList<>();
 
-    private final SimpleBroadcastReceiver mReceiver = new SimpleBroadcastReceiver(this::onIntent);
+    // We will register broadcast receiver on main thread to ensure not missing changes on
+    // TARGET_OVERLAY_PACKAGE and ACTION_OVERLAY_CHANGED.
+    private final SimpleBroadcastReceiver mReceiver =
+            new SimpleBroadcastReceiver(MAIN_EXECUTOR, this::onIntent);
 
     private Info mInfo;
     private boolean mDestroyed = false;
@@ -132,11 +135,11 @@
             mWindowContext.registerComponentCallbacks(this);
         } else {
             mWindowContext = null;
-            mReceiver.registerAsync(mContext, ACTION_CONFIGURATION_CHANGED);
+            mReceiver.register(mContext, ACTION_CONFIGURATION_CHANGED);
         }
 
         // Initialize navigation mode change listener
-        mReceiver.registerPkgActionsAsync(mContext, TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED);
+        mReceiver.registerPkgActions(mContext, TARGET_OVERLAY_PACKAGE, ACTION_OVERLAY_CHANGED);
 
         WindowManagerProxy wmProxy = WindowManagerProxy.INSTANCE.get(context);
         Context displayInfoContext = getDisplayInfoContext(display);
@@ -223,7 +226,7 @@
         } else {
             // TODO: unregister broadcast receiver
         }
-        mReceiver.unregisterReceiverSafelyAsync(mContext);
+        mReceiver.unregisterReceiverSafely(mContext);
     }
 
     /**
diff --git a/src/com/android/launcher3/util/LockedUserState.kt b/src/com/android/launcher3/util/LockedUserState.kt
index 2737249..10559f3 100644
--- a/src/com/android/launcher3/util/LockedUserState.kt
+++ b/src/com/android/launcher3/util/LockedUserState.kt
@@ -20,21 +20,28 @@
 import android.os.Process
 import android.os.UserManager
 import androidx.annotation.VisibleForTesting
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
 
 class LockedUserState(private val mContext: Context) : SafeCloseable {
     val isUserUnlockedAtLauncherStartup: Boolean
-    var isUserUnlocked: Boolean
-        private set
+    var isUserUnlocked = false
+        private set(value) {
+            field = value
+            if (value) {
+                notifyUserUnlocked()
+            }
+        }
 
     private val mUserUnlockedActions: RunnableList = RunnableList()
 
     @VisibleForTesting
-    val mUserUnlockedReceiver = SimpleBroadcastReceiver {
-        if (Intent.ACTION_USER_UNLOCKED == it.action) {
-            isUserUnlocked = true
-            notifyUserUnlocked()
+    val mUserUnlockedReceiver =
+        SimpleBroadcastReceiver(UI_HELPER_EXECUTOR) {
+            if (Intent.ACTION_USER_UNLOCKED == it.action) {
+                isUserUnlocked = true
+            }
         }
-    }
 
     init {
         // 1) when user reboots devices, launcher process starts at lock screen and both
@@ -43,26 +50,34 @@
         // yet isUserUnlockedAtLauncherStartup will remains as false.
         // 2) when launcher process restarts after user has unlocked screen, both variable are
         // init as true and will not change.
-        isUserUnlocked =
-            mContext
-                .getSystemService(UserManager::class.java)!!
-                .isUserUnlocked(Process.myUserHandle())
+        isUserUnlocked = checkIsUserUnlocked()
         isUserUnlockedAtLauncherStartup = isUserUnlocked
-        if (isUserUnlocked) {
-            notifyUserUnlocked()
-        } else {
-            mUserUnlockedReceiver.registerAsync(mContext, Intent.ACTION_USER_UNLOCKED)
+        if (!isUserUnlocked) {
+            mUserUnlockedReceiver.register(
+                mContext,
+                {
+                    // If user is unlocked while registering broadcast receiver, we should update
+                    // [isUserUnlocked], which will call [notifyUserUnlocked] in setter
+                    if (checkIsUserUnlocked()) {
+                        MAIN_EXECUTOR.execute { isUserUnlocked = true }
+                    }
+                },
+                Intent.ACTION_USER_UNLOCKED
+            )
         }
     }
 
+    private fun checkIsUserUnlocked() =
+        mContext.getSystemService(UserManager::class.java)!!.isUserUnlocked(Process.myUserHandle())
+
     private fun notifyUserUnlocked() {
         mUserUnlockedActions.executeAllAndDestroy()
-        mUserUnlockedReceiver.unregisterReceiverSafelyAsync(mContext)
+        mUserUnlockedReceiver.unregisterReceiverSafely(mContext)
     }
 
     /** Stops the receiver from listening for ACTION_USER_UNLOCK broadcasts. */
     override fun close() {
-        mUserUnlockedReceiver.unregisterReceiverSafelyAsync(mContext)
+        mUserUnlockedReceiver.unregisterReceiverSafely(mContext)
     }
 
     /**
diff --git a/src/com/android/launcher3/util/ScreenOnTracker.java b/src/com/android/launcher3/util/ScreenOnTracker.java
index c1d192c..12eff61 100644
--- a/src/com/android/launcher3/util/ScreenOnTracker.java
+++ b/src/com/android/launcher3/util/ScreenOnTracker.java
@@ -19,6 +19,8 @@
 import static android.content.Intent.ACTION_SCREEN_ON;
 import static android.content.Intent.ACTION_USER_PRESENT;
 
+import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
+
 import android.content.Context;
 import android.content.Intent;
 
@@ -32,7 +34,8 @@
     public static final MainThreadInitializedObject<ScreenOnTracker> INSTANCE =
             new MainThreadInitializedObject<>(ScreenOnTracker::new);
 
-    private final SimpleBroadcastReceiver mReceiver = new SimpleBroadcastReceiver(this::onReceive);
+    private final SimpleBroadcastReceiver mReceiver =
+            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::onReceive);
     private final CopyOnWriteArrayList<ScreenOnListener> mListeners = new CopyOnWriteArrayList<>();
 
     private final Context mContext;
@@ -42,12 +45,12 @@
         // Assume that the screen is on to begin with
         mContext = context;
         mIsScreenOn = true;
-        mReceiver.registerAsync(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
+        mReceiver.register(context, ACTION_SCREEN_ON, ACTION_SCREEN_OFF, ACTION_USER_PRESENT);
     }
 
     @Override
     public void close() {
-        mReceiver.unregisterReceiverSafelyAsync(mContext);
+        mReceiver.unregisterReceiverSafely(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 5f39cce..539a7cb 100644
--- a/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
+++ b/src/com/android/launcher3/util/SimpleBroadcastReceiver.java
@@ -15,21 +15,17 @@
  */
 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.Handler;
 import android.os.Looper;
 import android.os.PatternMatcher;
 import android.text.TextUtils;
 
+import androidx.annotation.AnyThread;
 import androidx.annotation.Nullable;
-import androidx.annotation.UiThread;
-import androidx.annotation.WorkerThread;
-
-import com.android.launcher3.BuildConfig;
 
 import java.util.function.Consumer;
 
@@ -37,8 +33,16 @@
 
     private final Consumer<Intent> mIntentConsumer;
 
-    public SimpleBroadcastReceiver(Consumer<Intent> intentConsumer) {
+    // Handler to register/unregister broadcast receiver
+    private final Handler mHandler;
+
+    public SimpleBroadcastReceiver(LooperExecutor looperExecutor, Consumer<Intent> intentConsumer) {
+        this(looperExecutor.getHandler(), intentConsumer);
+    }
+
+    public SimpleBroadcastReceiver(Handler handler, Consumer<Intent> intentConsumer) {
         mIntentConsumer = intentConsumer;
+        mHandler = handler;
     }
 
     @Override
@@ -46,55 +50,104 @@
         mIntentConsumer.accept(intent);
     }
 
-    /** 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));
+    /** Calls {@link #register(Context, Runnable, String...)} with null completionCallback. */
+    @AnyThread
+    public void register(Context context, String... actions) {
+        register(context, null, actions);
     }
 
-    /** Helper method to register multiple actions. Caller should be on main thread. */
-    @WorkerThread
-    public void registerSync(Context context, String... actions) {
-        assertOnBgThread();
+    /**
+     * Calls {@link #register(Context, Runnable, int, String...)} with null completionCallback.
+     */
+    @AnyThread
+    public void register(Context context, int flags, String... actions) {
+        register(context, null, flags, actions);
+    }
+
+    /**
+     * Register broadcast receiver. If this method is called on the same looper with mHandler's
+     * looper, then register will be called synchronously. Otherwise asynchronously. This ensures
+     * register happens on {@link #mHandler}'s looper.
+     *
+     * @param completionCallback callback that will be triggered after registration is completed,
+     *                           caller usually pass this callback to check if states has changed
+     *                           while registerReceiver() is executed on a binder call.
+     */
+    @AnyThread
+    public void register(
+            Context context, @Nullable Runnable completionCallback, String... actions) {
+        if (Looper.myLooper() == mHandler.getLooper()) {
+            registerInternal(context, completionCallback, actions);
+        } else {
+            mHandler.post(() -> registerInternal(context, completionCallback, actions));
+        }
+    }
+
+    /** Register broadcast receiver and run completion callback if passed. */
+    @AnyThread
+    private void registerInternal(
+            Context context, @Nullable Runnable completionCallback, String... actions) {
         context.registerReceiver(this, getFilter(actions));
+        if (completionCallback != null) {
+            completionCallback.run();
+        }
     }
 
     /**
-     * Helper method to register multiple actions associated with a action. Caller should be from
-     * main thread.
+     * Same as {@link #register(Context, Runnable, String...)} above but with additional flags
+     * params.
      */
-    @UiThread
-    public void registerPkgActionsAsync(Context context, @Nullable String pkg, String... actions) {
-        assertOnMainThread();
-        UI_HELPER_EXECUTOR.execute(() -> registerPkgActionsSync(context, pkg, actions));
+    @AnyThread
+    public void register(
+            Context context, @Nullable Runnable completionCallback, int flags, String... actions) {
+        if (Looper.myLooper() == mHandler.getLooper()) {
+            registerInternal(context, completionCallback, flags, actions);
+        } else {
+            mHandler.post(() -> registerInternal(context, completionCallback, flags, actions));
+        }
+    }
+
+    /** Register broadcast receiver and run completion callback if passed. */
+    @AnyThread
+    private void registerInternal(
+            Context context, @Nullable Runnable completionCallback, int flags, String... actions) {
+        context.registerReceiver(this, getFilter(actions), flags);
+        if (completionCallback != null) {
+            completionCallback.run();
+        }
+    }
+
+    /** Same as {@link #register(Context, Runnable, String...)} above but with pkg name. */
+    @AnyThread
+    public void registerPkgActions(Context context, @Nullable String pkg, String... actions) {
+        if (Looper.myLooper() == mHandler.getLooper()) {
+            context.registerReceiver(this, getPackageFilter(pkg, actions));
+        } else {
+            mHandler.post(() -> {
+                context.registerReceiver(this, getPackageFilter(pkg, actions));
+            });
+        }
     }
 
     /**
-     * Helper method to register multiple actions associated with a action. Caller should be from
-     * bg thread.
+     * Unregister broadcast receiver. If this method is called on the same looper with mHandler's
+     * looper, then unregister will be called synchronously. Otherwise asynchronously. This ensures
+     * unregister happens on {@link #mHandler}'s looper.
      */
-    @WorkerThread
-    public void registerPkgActionsSync(Context context, @Nullable String pkg, String... actions) {
-        assertOnBgThread();
-        context.registerReceiver(this, getPackageFilter(pkg, actions));
+    @AnyThread
+    public void unregisterReceiverSafely(Context context) {
+        if (Looper.myLooper() == mHandler.getLooper()) {
+            unregisterReceiverSafelyInternal(context);
+        } else {
+            mHandler.post(() -> {
+                unregisterReceiverSafelyInternal(context);
+            });
+        }
     }
 
-    /**
-     * 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();
+    /** Unregister broadcast receiver ignoring any errors. */
+    @AnyThread
+    private void unregisterReceiverSafelyInternal(Context context) {
         try {
             context.unregisterReceiver(this);
         } catch (IllegalArgumentException e) {
@@ -121,20 +174,4 @@
         }
         return filter;
     }
-
-    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 a2277a0..f8cbe0d 100644
--- a/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
+++ b/src/com/android/launcher3/util/WallpaperOffsetInterpolator.java
@@ -32,7 +32,7 @@
     private static final int MIN_PARALLAX_PAGE_SPAN = 4;
 
     private final SimpleBroadcastReceiver mWallpaperChangeReceiver =
-            new SimpleBroadcastReceiver(i -> onWallpaperChanged());
+            new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, i -> onWallpaperChanged());
     private final Workspace<?> mWorkspace;
     private final boolean mIsRtl;
     private final Handler mHandler;
@@ -198,10 +198,10 @@
     public void setWindowToken(IBinder token) {
         mWindowToken = token;
         if (mWindowToken == null && mRegistered) {
-            mWallpaperChangeReceiver.unregisterReceiverSafelyAsync(mWorkspace.getContext());
+            mWallpaperChangeReceiver.unregisterReceiverSafely(mWorkspace.getContext());
             mRegistered = false;
         } else if (mWindowToken != null && !mRegistered) {
-            mWallpaperChangeReceiver.registerAsync(
+            mWallpaperChangeReceiver.register(
                     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 6e01f9e..3d253b4 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -22,6 +22,7 @@
 import static com.android.launcher3.testing.shared.TestProtocol.ICON_MISSING;
 import static com.android.launcher3.testing.shared.TestProtocol.WIDGET_CONFIG_NULL_EXTRA_INTENT;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
+import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -238,9 +239,9 @@
     protected void clearPackageData(String pkg) throws IOException, InterruptedException {
         final CountDownLatch count = new CountDownLatch(2);
         final SimpleBroadcastReceiver broadcastReceiver =
-                new SimpleBroadcastReceiver(i -> count.countDown());
+                new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, i -> count.countDown());
         // We OK to make binder calls on main thread in test.
-        broadcastReceiver.registerPkgActionsSync(mTargetContext, pkg,
+        broadcastReceiver.registerPkgActions(mTargetContext, pkg,
                 Intent.ACTION_PACKAGE_RESTARTED, Intent.ACTION_PACKAGE_DATA_CLEARED);
 
         mDevice.executeShellCommand("pm clear " + pkg);
diff --git a/tests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt b/tests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt
new file mode 100644
index 0000000..1de99c5
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/SimpleBroadcastReceiverTest.kt
@@ -0,0 +1,158 @@
+/*
+ * Copyright (C) 2024 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.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.Handler
+import android.os.Looper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
+import com.google.common.truth.Truth.assertThat
+import java.util.function.Consumer
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.eq
+import org.mockito.ArgumentMatchers.same
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.verify
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class SimpleBroadcastReceiverTest {
+
+    private lateinit var underTest: SimpleBroadcastReceiver
+
+    @Mock private lateinit var intentConsumer: Consumer<Intent>
+    @Mock private lateinit var context: Context
+    @Mock private lateinit var completionRunnable: Runnable
+    @Captor private lateinit var intentFilterCaptor: ArgumentCaptor<IntentFilter>
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        underTest = SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, intentConsumer)
+        if (Looper.getMainLooper() == null) {
+            Looper.prepareMainLooper()
+        }
+    }
+
+    @Test
+    fun async_register() {
+        underTest.register(context, "test_action_1", "test_action_2")
+        awaitTasksCompleted()
+
+        verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture())
+        val intentFilter = intentFilterCaptor.value
+        assertThat(intentFilter.countActions()).isEqualTo(2)
+        assertThat(intentFilter.getAction(0)).isEqualTo("test_action_1")
+        assertThat(intentFilter.getAction(1)).isEqualTo("test_action_2")
+    }
+
+    @Test
+    fun async_register_withCompletionRunnable() {
+        underTest.register(context, completionRunnable, "test_action_1", "test_action_2")
+        awaitTasksCompleted()
+
+        verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture())
+        verify(completionRunnable).run()
+        val intentFilter = intentFilterCaptor.value
+        assertThat(intentFilter.countActions()).isEqualTo(2)
+        assertThat(intentFilter.getAction(0)).isEqualTo("test_action_1")
+        assertThat(intentFilter.getAction(1)).isEqualTo("test_action_2")
+    }
+
+    @Test
+    fun async_register_withCompletionRunnable_and_flag() {
+        underTest.register(context, completionRunnable, 1, "test_action_1", "test_action_2")
+        awaitTasksCompleted()
+
+        verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture(), eq(1))
+        verify(completionRunnable).run()
+        val intentFilter = intentFilterCaptor.value
+        assertThat(intentFilter.countActions()).isEqualTo(2)
+        assertThat(intentFilter.getAction(0)).isEqualTo("test_action_1")
+        assertThat(intentFilter.getAction(1)).isEqualTo("test_action_2")
+    }
+
+    @Test
+    fun async_register_with_package() {
+        underTest.registerPkgActions(context, "pkg", "test_action_1", "test_action_2")
+
+        awaitTasksCompleted()
+        verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture())
+        val intentFilter = intentFilterCaptor.value
+        assertThat(intentFilter.getDataScheme(0)).isEqualTo("package")
+        assertThat(intentFilter.getDataSchemeSpecificPart(0).path).isEqualTo("pkg")
+        assertThat(intentFilter.countActions()).isEqualTo(2)
+        assertThat(intentFilter.getAction(0)).isEqualTo("test_action_1")
+        assertThat(intentFilter.getAction(1)).isEqualTo("test_action_2")
+    }
+
+    @Test
+    fun sync_register_withCompletionRunnable_and_flag() {
+        underTest = SimpleBroadcastReceiver(Handler(Looper.getMainLooper()), intentConsumer)
+
+        underTest.register(context, completionRunnable, 1, "test_action_1", "test_action_2")
+
+        verify(context).registerReceiver(same(underTest), intentFilterCaptor.capture(), eq(1))
+        verify(completionRunnable).run()
+        val intentFilter = intentFilterCaptor.value
+        assertThat(intentFilter.countActions()).isEqualTo(2)
+        assertThat(intentFilter.getAction(0)).isEqualTo("test_action_1")
+        assertThat(intentFilter.getAction(1)).isEqualTo("test_action_2")
+    }
+
+    @Test
+    fun async_unregister() {
+        underTest.unregisterReceiverSafely(context)
+
+        awaitTasksCompleted()
+        verify(context).unregisterReceiver(same(underTest))
+    }
+
+    @Test
+    fun sync_unregister() {
+        underTest = SimpleBroadcastReceiver(Handler(Looper.getMainLooper()), intentConsumer)
+
+        underTest.unregisterReceiverSafely(context)
+
+        verify(context).unregisterReceiver(same(underTest))
+    }
+
+    @Test
+    fun getPackageFilter() {
+        val intentFilter =
+            SimpleBroadcastReceiver.getPackageFilter("pkg", "test_action_1", "test_action_2")
+
+        assertThat(intentFilter.getDataScheme(0)).isEqualTo("package")
+        assertThat(intentFilter.getDataSchemeSpecificPart(0).path).isEqualTo("pkg")
+        assertThat(intentFilter.countActions()).isEqualTo(2)
+        assertThat(intentFilter.getAction(0)).isEqualTo("test_action_1")
+        assertThat(intentFilter.getAction(1)).isEqualTo("test_action_2")
+    }
+
+    private fun awaitTasksCompleted() {
+        UI_HELPER_EXECUTOR.submit<Any> { null }.get()
+    }
+}