[Launcher Jank] Improve SimpleBroadcastReceiver.java

Bug: 348649441
Flag: NONE - jank fix
Test: manual - presubmit
Change-Id: I17bd7e4d7f0640522476e94edae691f983835f62
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;