Merge "Prevent ConcurrentModificationException" into main
diff --git a/packages/SystemUI/shared/src/com/android/systemui/statusbar/policy/CallbackController.java b/packages/SystemUI/shared/src/com/android/systemui/statusbar/policy/CallbackController.java
index 047ff75..9f82201 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/statusbar/policy/CallbackController.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/statusbar/policy/CallbackController.java
@@ -21,6 +21,11 @@
 import androidx.lifecycle.LifecycleEventObserver;
 import androidx.lifecycle.LifecycleOwner;
 
+/**
+ * Implementation of the collection used and thread guarantees are left to the discretion of the
+ * client. Consider using {@link com.android.systemui.util.ListenerSet} to prevent concurrent
+ * modification exceptions.
+ */
 public interface CallbackController<T> {
 
     /** Add a callback */
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ReduceBrightColorsController.java b/packages/SystemUI/src/com/android/systemui/qs/ReduceBrightColorsController.java
index 36dc743..a01d658 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ReduceBrightColorsController.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/ReduceBrightColorsController.java
@@ -67,9 +67,7 @@
                 synchronized (mListeners) {
                     if (setting != null && mListeners.size() != 0) {
                         if (setting.equals(Settings.Secure.REDUCE_BRIGHT_COLORS_ACTIVATED)) {
-                            for (Listener listener : mListeners) {
-                                listener.onActivated(mManager.isReduceBrightColorsActivated());
-                            }
+                            dispatchOnActivated(mManager.isReduceBrightColorsActivated());
                         }
                     }
                 }
@@ -125,6 +123,13 @@
         mManager.setReduceBrightColorsActivated(activated);
     }
 
+    private void dispatchOnActivated(boolean activated) {
+        ArrayList<Listener> copy = new ArrayList<>(mListeners);
+        for (Listener l : copy) {
+            l.onActivated(activated);
+        }
+    }
+
     /**
      * Listener invoked whenever the Reduce Bright Colors settings are changed.
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
index abdf827..56ea00c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
@@ -104,7 +104,8 @@
     }
 
     private void notifyManagedProfileRemoved() {
-        for (Callback callback : mCallbacks) {
+        ArrayList<Callback> copy = new ArrayList<>(mCallbacks);
+        for (Callback callback : copy) {
             callback.onManagedProfileRemoved();
         }
     }
@@ -148,7 +149,8 @@
         @Override
         public void onUserChanged(int newUser, @NonNull Context userContext) {
             reloadManagedProfiles();
-            for (Callback callback : mCallbacks) {
+            ArrayList<Callback> copy = new ArrayList<>(mCallbacks);
+            for (Callback callback : copy) {
                 callback.onManagedProfileChanged();
             }
         }
@@ -156,7 +158,8 @@
         @Override
         public void onProfilesChanged(@NonNull List<UserInfo> profiles) {
             reloadManagedProfiles();
-            for (Callback callback : mCallbacks) {
+            ArrayList<Callback> copy = new ArrayList<>(mCallbacks);
+            for (Callback callback : copy) {
                 callback.onManagedProfileChanged();
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java
index 41ed76d..45078e3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java
@@ -20,6 +20,7 @@
 import static android.os.BatteryManager.CHARGING_POLICY_DEFAULT;
 import static android.os.BatteryManager.EXTRA_CHARGING_STATUS;
 import static android.os.BatteryManager.EXTRA_PRESENT;
+
 import static com.android.settingslib.fuelgauge.BatterySaverLogging.SAVER_ENABLED_QS;
 import static com.android.systemui.util.DumpUtilsKt.asIndenting;
 
@@ -61,6 +62,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 import javax.annotation.concurrent.GuardedBy;
 
@@ -448,50 +450,38 @@
         firePowerSaveChanged();
     }
 
+    protected final void dispatchSafeChange(Consumer<BatteryStateChangeCallback> action) {
+        ArrayList<BatteryStateChangeCallback> copy;
+        synchronized (mChangeCallbacks) {
+            copy = new ArrayList<>(mChangeCallbacks);
+        }
+        final int n = copy.size();
+        for (int i = 0; i < n; i++) {
+            action.accept(copy.get(i));
+        }
+    }
+
     protected void fireBatteryLevelChanged() {
         mLogger.logBatteryLevelChangedCallback(mLevel, mPluggedIn, mCharging);
-        synchronized (mChangeCallbacks) {
-            final int N = mChangeCallbacks.size();
-            for (int i = 0; i < N; i++) {
-                mChangeCallbacks.get(i).onBatteryLevelChanged(mLevel, mPluggedIn, mCharging);
-            }
-        }
+        dispatchSafeChange(
+                (callback) -> callback.onBatteryLevelChanged(mLevel, mPluggedIn, mCharging));
     }
 
     private void fireBatteryUnknownStateChanged() {
-        synchronized (mChangeCallbacks) {
-            final int n = mChangeCallbacks.size();
-            for (int i = 0; i < n; i++) {
-                mChangeCallbacks.get(i).onBatteryUnknownStateChanged(mStateUnknown);
-            }
-        }
+        dispatchSafeChange((callback) -> callback.onBatteryUnknownStateChanged(mStateUnknown));
     }
 
     private void firePowerSaveChanged() {
-        synchronized (mChangeCallbacks) {
-            final int N = mChangeCallbacks.size();
-            for (int i = 0; i < N; i++) {
-                mChangeCallbacks.get(i).onPowerSaveChanged(mPowerSave);
-            }
-        }
+        dispatchSafeChange((callback) -> callback.onPowerSaveChanged(mPowerSave));
     }
 
     private void fireIsBatteryDefenderChanged() {
-        synchronized (mChangeCallbacks) {
-            final int n = mChangeCallbacks.size();
-            for (int i = 0; i < n; i++) {
-                mChangeCallbacks.get(i).onIsBatteryDefenderChanged(mIsBatteryDefender);
-            }
-        }
+        dispatchSafeChange((callback) -> callback.onIsBatteryDefenderChanged(mIsBatteryDefender));
     }
 
     private void fireIsIncompatibleChargingChanged() {
-        synchronized (mChangeCallbacks) {
-            final int n = mChangeCallbacks.size();
-            for (int i = 0; i < n; i++) {
-                mChangeCallbacks.get(i).onIsIncompatibleChargingChanged(mIsIncompatibleCharging);
-            }
-        }
+        dispatchSafeChange(
+                (callback) -> callback.onIsIncompatibleChargingChanged(mIsIncompatibleCharging));
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BluetoothControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BluetoothControllerImpl.java
index 53b343c..fc2f6e9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BluetoothControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BluetoothControllerImpl.java
@@ -436,6 +436,8 @@
     @Override
     public void onServiceDisconnected() {}
 
+    // IMPORTANT: This handler guarantees that any operations on the list of callbacks is
+    // sequential, so no concurrent exceptions
     private final class H extends Handler {
         private final ArrayList<BluetoothController.Callback> mCallbacks = new ArrayList<>();
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/CastControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/CastControllerImpl.java
index b06ebe9..149c8fa 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/CastControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/CastControllerImpl.java
@@ -35,9 +35,9 @@
 import androidx.annotation.VisibleForTesting;
 
 import com.android.internal.annotations.GuardedBy;
-import com.android.systemui.res.R;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.res.R;
 import com.android.systemui.util.Utils;
 
 import java.io.PrintWriter;
@@ -291,11 +291,12 @@
 
     @VisibleForTesting
     void fireOnCastDevicesChanged() {
+        final ArrayList<Callback> callbacks;
         synchronized (mCallbacks) {
-            for (Callback callback : mCallbacks) {
-                fireOnCastDevicesChanged(callback);
-            }
-
+            callbacks = new ArrayList<>(mCallbacks);
+        }
+        for (Callback callback : callbacks) {
+            fireOnCastDevicesChanged(callback);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/DataSaverControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/DataSaverControllerImpl.java
index 8207012..6319781 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/DataSaverControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/DataSaverControllerImpl.java
@@ -36,10 +36,12 @@
     }
 
     private void handleRestrictBackgroundChanged(boolean isDataSaving) {
+        ArrayList<DataSaverController.Listener> copy;
         synchronized (mListeners) {
-            for (int i = 0; i < mListeners.size(); i++) {
-                mListeners.get(i).onDataSaverChanged(isDataSaving);
-            }
+            copy = new ArrayList<>(mListeners);
+        }
+        for (int i = 0; i < copy.size(); i++) {
+            copy.get(i).onDataSaverChanged(isDataSaving);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/FlashlightControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/FlashlightControllerImpl.java
index 5dcafb3..b98eff8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/FlashlightControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/FlashlightControllerImpl.java
@@ -202,10 +202,12 @@
 
     private void dispatchListeners(int message, boolean argument) {
         synchronized (mListeners) {
-            final int N = mListeners.size();
+            final ArrayList<WeakReference<FlashlightController.FlashlightListener>> copy =
+                    new ArrayList<>(mListeners);
+            final int n = copy.size();
             boolean cleanup = false;
-            for (int i = 0; i < N; i++) {
-                FlashlightListener l = mListeners.get(i).get();
+            for (int i = 0; i < n; i++) {
+                FlashlightListener l = copy.get(i).get();
                 if (l != null) {
                     if (message == DISPATCH_ERROR) {
                         l.onFlashlightError();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/IndividualSensorPrivacyControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/IndividualSensorPrivacyControllerImpl.java
index fffd839..87dfc99 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/IndividualSensorPrivacyControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/IndividualSensorPrivacyControllerImpl.java
@@ -119,7 +119,8 @@
             mHardwareToggleState.put(sensor, enabled);
         }
 
-        for (Callback callback : mCallbacks) {
+        Set<Callback> copy = new ArraySet<>(mCallbacks);
+        for (Callback callback : copy) {
             callback.onSensorBlockedChanged(sensor, isSensorBlocked(sensor));
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/LocationControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/LocationControllerImpl.java
index e5f72eb..9eee5d0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/LocationControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/LocationControllerImpl.java
@@ -356,6 +356,8 @@
         updateActiveLocationRequests();
     }
 
+    // IMPORTANT: This handler guarantees that any operations on the list of callbacks is
+    // sequential, so no concurrent exceptions
     private final class H extends Handler {
         private static final int MSG_LOCATION_SETTINGS_CHANGED = 1;
         private static final int MSG_LOCATION_ACTIVE_CHANGED = 2;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/NextAlarmControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/NextAlarmControllerImpl.java
index 63b9ff9..b7d8ee3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/NextAlarmControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/NextAlarmControllerImpl.java
@@ -124,9 +124,10 @@
     }
 
     private void fireNextAlarmChanged() {
-        int n = mChangeCallbacks.size();
+        ArrayList<NextAlarmChangeCallback> copy = new ArrayList<>(mChangeCallbacks);
+        int n = copy.size();
         for (int i = 0; i < n; i++) {
-            mChangeCallbacks.get(i).onNextAlarmChanged(mNextAlarm);
+            copy.get(i).onNextAlarmChanged(mNextAlarm);
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SafetyController.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SafetyController.java
index f3d183c..0176abd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SafetyController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SafetyController.java
@@ -100,10 +100,13 @@
     }
 
     private void handleSafetyCenterEnableChange() {
+        final ArrayList<SafetyController.Listener> copy;
         synchronized (mListeners) {
-            for (int i = 0; i < mListeners.size(); i++) {
-                mListeners.get(i).onSafetyCenterEnableChanged(mSafetyCenterEnabled);
-            }
+            copy = new ArrayList<>(mListeners);
+        }
+        final int n = copy.size();
+        for (int i = 0; i < n; i++) {
+            copy.get(i).onSafetyCenterEnableChanged(mSafetyCenterEnabled);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
index 4a4d4e1..5d69f36 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
@@ -50,12 +50,12 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.net.LegacyVpnInfo;
 import com.android.internal.net.VpnConfig;
-import com.android.systemui.res.R;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.res.R;
 import com.android.systemui.settings.UserTracker;
 
 import org.xmlpull.v1.XmlPullParserException;
@@ -429,10 +429,12 @@
     }
 
     private void fireCallbacks() {
+        final ArrayList<SecurityControllerCallback> copy;
         synchronized (mCallbacks) {
-            for (SecurityControllerCallback callback : mCallbacks) {
-                callback.onStateChanged();
-            }
+            copy = new ArrayList<>(mCallbacks);
+        }
+        for (SecurityControllerCallback callback : copy) {
+            callback.onStateChanged();
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java
index 66bf527..df210b0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java
@@ -48,12 +48,12 @@
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.settings.UserTracker;
-import com.android.systemui.util.Utils;
 import com.android.systemui.util.settings.GlobalSettings;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 import javax.inject.Inject;
 
@@ -243,46 +243,43 @@
     }
 
     private void fireNextAlarmChanged() {
-        synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onNextAlarmChanged());
-        }
+        fireSafeChange(Callback::onNextAlarmChanged);
     }
 
     private void fireEffectsSuppressorChanged() {
-        synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onEffectsSupressorChanged());
-        }
+        fireSafeChange(Callback::onEffectsSupressorChanged);
     }
 
     private void fireZenChanged(int zen) {
-        synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onZenChanged(zen));
-        }
+        fireSafeChange(c -> c.onZenChanged(zen));
     }
 
     private void fireZenAvailableChanged(boolean available) {
-        synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onZenAvailableChanged(available));
-        }
+        fireSafeChange(c -> c.onZenAvailableChanged(available));
     }
 
     private void fireManualRuleChanged(ZenRule rule) {
-        synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onManualRuleChanged(rule));
-        }
+        fireSafeChange(c -> c.onManualRuleChanged(rule));
     }
 
     private void fireConsolidatedPolicyChanged(NotificationManager.Policy policy) {
+        fireSafeChange(c -> c.onConsolidatedPolicyChanged(policy));
+    }
+
+    private void fireSafeChange(Consumer<Callback> action) {
+        final ArrayList<Callback> copy;
         synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onConsolidatedPolicyChanged(policy));
+            copy = new ArrayList<>(mCallbacks);
+        }
+        final int n = copy.size();
+        for (int i = 0; i < n; i++) {
+            action.accept(copy.get(i));
         }
     }
 
     @VisibleForTesting
     protected void fireConfigChanged(ZenModeConfig config) {
-        synchronized (mCallbacksLock) {
-            Utils.safeForeach(mCallbacks, c -> c.onConfigChanged(config));
-        }
+        fireSafeChange(c -> c.onConfigChanged(config));
     }
 
     @VisibleForTesting
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
index 7eba3b46..c44c178 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
@@ -22,11 +22,15 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.time.FakeSystemClock
 import junit.framework.Assert
 import org.junit.Before
 import org.junit.Test
 import org.mockito.Mock
+import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
 
@@ -72,6 +76,34 @@
         Assert.assertEquals(true, controller.hasActiveProfile())
     }
 
+    @Test
+    fun callbackRemovedWhileDispatching_doesntCrash() {
+        var remove = false
+        val callback =
+            object : ManagedProfileController.Callback {
+                override fun onManagedProfileChanged() {
+                    if (remove) {
+                        controller.removeCallback(this)
+                    }
+                }
+
+                override fun onManagedProfileRemoved() {
+                    if (remove) {
+                        controller.removeCallback(this)
+                    }
+                }
+            }
+        controller.addCallback(callback)
+        controller.addCallback(TestCallback)
+
+        remove = true
+        setupWorkingProfile(1)
+
+        val captor = argumentCaptor<UserTracker.Callback>()
+        verify(userTracker).addCallback(capture(captor), any())
+        captor.value.onProfilesChanged(userManager.getEnabledProfiles(1))
+    }
+
     private fun setupWorkingProfile(userId: Int) {
         `when`(userManager.getEnabledProfiles(userId))
             .thenReturn(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java
index 2f79955..a5c766d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java
@@ -17,10 +17,12 @@
 package com.android.systemui.statusbar.policy;
 
 import static android.os.BatteryManager.EXTRA_PRESENT;
+
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.inOrder;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.staticMockMarker;
 import static com.android.settingslib.fuelgauge.BatterySaverLogging.SAVER_ENABLED_QS;
+
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -59,6 +61,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
@@ -286,6 +289,24 @@
         Assert.assertFalse(mBatteryController.isIncompatibleCharging());
     }
 
+    @Test
+    public void callbackRemovedWhileDispatching_doesntCrash() {
+        final AtomicBoolean remove = new AtomicBoolean(false);
+        BatteryStateChangeCallback callback = new BatteryStateChangeCallback() {
+            @Override
+            public void onBatteryLevelChanged(int level, boolean pluggedIn, boolean charging) {
+                if (remove.get()) {
+                    mBatteryController.removeCallback(this);
+                }
+            }
+        };
+        mBatteryController.addCallback(callback);
+        // Add another callback so the iteration continues
+        mBatteryController.addCallback(new BatteryStateChangeCallback() {});
+        remove.set(true);
+        mBatteryController.fireBatteryLevelChanged();
+    }
+
     private void setupIncompatibleCharging() {
         final List<UsbPort> usbPorts = new ArrayList<>();
         usbPorts.add(mUsbPort);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java
index b8e4306..68c1b8d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java
@@ -127,6 +127,25 @@
         }
     }
 
+    /** Regression test for b/317700495 */
+    @Test
+    public void removeCallbackWhileIterating_doesntCrash() {
+        final AtomicBoolean remove = new AtomicBoolean(false);
+        Callback callback = new Callback() {
+            @Override
+            public void onCastDevicesChanged() {
+                if (remove.get()) {
+                    mController.removeCallback(this);
+                }
+            }
+        };
+        mController.addCallback(callback);
+        // Add another callback so the iteration continues
+        mController.addCallback(() -> {});
+        remove.set(true);
+        mController.fireOnCastDevicesChanged();
+    }
+
     @Test
     public void hasConnectedCastDevice_connected() {
         CastController.CastDevice castDevice = new CastController.CastDevice();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/FlashlightControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/FlashlightControllerImplTest.kt
index db0029a..777fa28 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/FlashlightControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/FlashlightControllerImplTest.kt
@@ -26,6 +26,8 @@
 import com.android.systemui.broadcast.BroadcastSender
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.settings.FakeSettings
 import com.android.systemui.util.time.FakeSystemClock
 import java.util.concurrent.Executor
@@ -128,11 +130,42 @@
         verify(cameraManager).setTorchMode(id, enable)
     }
 
+    @Test
+    fun testCallbackRemovedWhileDispatching_doesntCrash() {
+        injectCamera()
+        var remove = false
+        val callback = object : FlashlightController.FlashlightListener {
+            override fun onFlashlightChanged(enabled: Boolean) {
+                if (remove) {
+                    controller.removeCallback(this)
+                }
+            }
+
+            override fun onFlashlightError() {}
+
+            override fun onFlashlightAvailabilityChanged(available: Boolean) {}
+        }
+        controller.addCallback(callback)
+        controller.addCallback(object : FlashlightController.FlashlightListener {
+            override fun onFlashlightChanged(enabled: Boolean) {}
+
+            override fun onFlashlightError() {}
+
+            override fun onFlashlightAvailabilityChanged(available: Boolean) {}
+        })
+        backgroundExecutor.runAllReady()
+
+        val captor = argumentCaptor<CameraManager.TorchCallback>()
+        verify(cameraManager).registerTorchCallback(any(), capture(captor))
+        remove = true
+        captor.value.onTorchModeChanged(ID, true)
+    }
+
     private fun injectCamera(
         flash: Boolean = true,
         facing: Int = CameraCharacteristics.LENS_FACING_BACK
     ): String {
-        val cameraID = "ID"
+        val cameraID = ID
         val camera = CameraCharacteristics(CameraMetadataNative().apply {
             set(CameraCharacteristics.FLASH_INFO_AVAILABLE, flash)
             set(CameraCharacteristics.LENS_FACING, facing)
@@ -141,4 +174,8 @@
         `when`(cameraManager.getCameraCharacteristics(cameraID)).thenReturn(camera)
         return cameraID
     }
+
+    companion object {
+        private const val ID = "ID"
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt
index 89989ce..b03edaf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt
@@ -27,6 +27,7 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.util.mockito.any
+import com.google.common.truth.Truth.assertThat
 import org.junit.Assert.assertEquals
 import org.junit.Before
 import org.junit.Test
@@ -130,4 +131,61 @@
         controller.mPermControllerChangeReceiver.onReceive(context, testIntent)
         verify(listener, never()).onSafetyCenterEnableChanged(true)
     }
+
+    @Test
+    fun listenerRemovedWhileDispatching_doesNotCrash() {
+        var remove = false
+        val callback = object : SafetyController.Listener {
+            override fun onSafetyCenterEnableChanged(isSafetyCenterEnabled: Boolean) {
+                if (remove) {
+                    controller.removeCallback(this)
+                }
+            }
+        }
+
+        controller.addCallback(callback)
+        controller.addCallback {}
+
+        remove = true
+
+        `when`(scm.isSafetyCenterEnabled).thenReturn(true)
+        val testIntent = Intent(Intent.ACTION_PACKAGE_CHANGED)
+        testIntent.data = Uri.parse("package:$TEST_PC_PKG")
+        controller.mPermControllerChangeReceiver.onReceive(context, testIntent)
+    }
+
+    @Test
+    fun listenerRemovedWhileDispatching_otherCallbacksCalled() {
+        var remove = false
+        var called = false
+
+        val callback1 = object : SafetyController.Listener {
+            override fun onSafetyCenterEnableChanged(isSafetyCenterEnabled: Boolean) {
+                if (remove) {
+                    controller.removeCallback(this)
+                }
+            }
+        }
+
+        val callback2 = object : SafetyController.Listener {
+            override fun onSafetyCenterEnableChanged(isSafetyCenterEnabled: Boolean) {
+                // When the first callback is removed, we track if this is called
+                if (remove) {
+                    called = true
+                }
+            }
+        }
+
+        controller.addCallback(callback1)
+        controller.addCallback(callback2)
+
+        remove = true
+
+        `when`(scm.isSafetyCenterEnabled).thenReturn(true)
+        val testIntent = Intent(Intent.ACTION_PACKAGE_CHANGED)
+        testIntent.data = Uri.parse("package:$TEST_PC_PKG")
+        controller.mPermControllerChangeReceiver.onReceive(context, testIntent)
+
+        assertThat(called).isTrue()
+    }
 }
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
index c35bc69..1dab84e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
@@ -63,6 +63,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -217,6 +218,28 @@
                 ), any(NetworkCallback.class));
     }
 
+    @Test
+    public void testRemoveCallbackWhileDispatch_doesntCrash() {
+        final AtomicBoolean remove = new AtomicBoolean(false);
+        SecurityController.SecurityControllerCallback callback =
+                new SecurityController.SecurityControllerCallback() {
+                    @Override
+                    public void onStateChanged() {
+                        if (remove.get()) {
+                            mSecurityController.removeCallback(this);
+                        }
+                    }
+                };
+        mSecurityController.addCallback(callback);
+        // Add another callback so the iteration continues
+        mSecurityController.addCallback(() -> {});
+        mBgExecutor.runAllReady();
+        remove.set(true);
+
+        mSecurityController.onUserSwitched(10);
+        mBgExecutor.runAllReady();
+    }
+
     /**
      * refresh CA certs by sending a user unlocked broadcast for the desired user
      */
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java
index 6825f65..f1a2c28 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java
@@ -47,6 +47,7 @@
 import org.mockito.MockitoAnnotations;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 @SmallTest
@@ -174,4 +175,26 @@
         }
 
     }
+
+    @Test
+    public void testCallbackRemovedWhileDispatching_doesntCrash() {
+        final AtomicBoolean remove = new AtomicBoolean(false);
+        mGlobalSettings.putInt(Settings.Global.ZEN_MODE, Settings.Global.ZEN_MODE_OFF);
+        TestableLooper.get(this).processAllMessages();
+        final ZenModeController.Callback callback = new ZenModeController.Callback() {
+            @Override
+            public void onZenChanged(int zen) {
+                if (remove.get()) {
+                    mController.removeCallback(this);
+                }
+            }
+        };
+        mController.addCallback(callback);
+        mController.addCallback(new ZenModeController.Callback() {});
+
+        remove.set(true);
+
+        mGlobalSettings.putInt(Settings.Global.ZEN_MODE, Settings.Global.ZEN_MODE_NO_INTERRUPTIONS);
+        TestableLooper.get(this).processAllMessages();
+    }
 }