Copy keys to iterate over so we don't concurrently modify the map.

Bug: 159923411
Test: it's very sporadic unfortunately
Change-Id: I28e938b944bb52136c7b332b80a0f46dd73b448b

2 files changed, 41 insertions, 8 deletions

diff --git a/services/core/java/com/android/server/notification/ShortcutHelper.java b/services/core/java/com/android/server/notification/ShortcutHelper.java
index ee02e3fa8140..94b690a4dfce 100644
--- a/services/core/java/com/android/server/notification/ShortcutHelper.java
+++ b/services/core/java/com/android/server/notification/ShortcutHelper.java
@@ -37,7 +37,9 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 /**
  * Helper for querying shortcuts.
@@ -249,8 +251,11 @@ public class ShortcutHelper {
                 if (!TextUtils.isEmpty(shortcutId)) {
                     packageBubbles.remove(shortcutId);
                 } else {
+                    // Copy the shortcut IDs to avoid a concurrent modification exception.
+                    final Set<String> shortcutIds = new HashSet<>(packageBubbles.keySet());
+
                     // Check if there was a matching entry
-                    for (String pkgShortcutId : packageBubbles.keySet()) {
+                    for (String pkgShortcutId : shortcutIds) {
                         String entryKey = packageBubbles.get(pkgShortcutId);
                         if (r.getKey().equals(entryKey)) {
                             // No longer has shortcut id so remove it
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java
index c700a090fa2e..eca71b69ec0b 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ShortcutHelperTest.java
@@ -45,6 +45,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
+import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
@@ -85,14 +86,19 @@ public class ShortcutHelperTest extends UiServiceTestCase {
 
         mShortcutHelper = new ShortcutHelper(
                 mLauncherApps, mShortcutListener, mShortcutServiceInternal);
-        when(mNr.getKey()).thenReturn(KEY);
-        when(mNr.getSbn()).thenReturn(mSbn);
         when(mSbn.getPackageName()).thenReturn(PKG);
-        when(mNr.getNotification()).thenReturn(mNotif);
-        when(mNr.getShortcutInfo()).thenReturn(mShortcutInfo);
         when(mShortcutInfo.getId()).thenReturn(SHORTCUT_ID);
         when(mNotif.getBubbleMetadata()).thenReturn(mBubbleMetadata);
         when(mBubbleMetadata.getShortcutId()).thenReturn(SHORTCUT_ID);
+
+        setUpMockNotificationRecord(mNr, KEY);
+    }
+
+    private void setUpMockNotificationRecord(NotificationRecord mockRecord, String key) {
+        when(mockRecord.getKey()).thenReturn(key);
+        when(mockRecord.getSbn()).thenReturn(mSbn);
+        when(mockRecord.getNotification()).thenReturn(mNotif);
+        when(mockRecord.getShortcutInfo()).thenReturn(mShortcutInfo);
     }
 
     private LauncherApps.Callback addShortcutBubbleAndVerifyListener() {
@@ -159,9 +165,31 @@ public class ShortcutHelperTest extends UiServiceTestCase {
         // First set it up to listen
         addShortcutBubbleAndVerifyListener();
 
-        // Clear out shortcutId
-        when(mNr.getShortcutInfo()).thenReturn(null);
-        mShortcutHelper.maybeListenForShortcutChangesForBubbles(mNr,
+        NotificationRecord validMock1 = Mockito.mock(NotificationRecord.class);
+        setUpMockNotificationRecord(validMock1, "KEY1");
+
+        NotificationRecord validMock2 = Mockito.mock(NotificationRecord.class);
+        setUpMockNotificationRecord(validMock2, "KEY2");
+
+        NotificationRecord validMock3 = Mockito.mock(NotificationRecord.class);
+        setUpMockNotificationRecord(validMock3, "KEY3");
+
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock1,
+                false /* removed */,
+                null /* handler */);
+
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock2,
+                false /* removed */,
+                null /* handler */);
+
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock3,
+                false /* removed */,
+                null /* handler */);
+
+        // Clear out shortcutId of the bubble in the middle, to double check that we don't hit a
+        // concurrent modification exception (removing the last bubble would sidestep that check).
+        when(validMock2.getShortcutInfo()).thenReturn(null);
+        mShortcutHelper.maybeListenForShortcutChangesForBubbles(validMock2,
                 false /* removed */,
                 null /* handler */);