diff --git a/src/com/android/launcher3/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 97e93a0..24d16d7 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -49,7 +49,6 @@
 import com.android.launcher3.graphics.IconPalette;
 import com.android.launcher3.graphics.PreloadIconDrawable;
 import com.android.launcher3.model.PackageItemInfo;
-import com.android.launcher3.popup.PopupContainerWithArrow;
 
 import java.text.NumberFormat;
 
@@ -573,18 +572,11 @@
 
     public void applyBadgeState(ItemInfo itemInfo, boolean animate) {
         if (mIcon instanceof FastBitmapDrawable) {
-            BadgeInfo badgeInfo = mLauncher.getPopupDataProvider().getBadgeInfoForItem(itemInfo);
-            BadgeRenderer badgeRenderer = mLauncher.getDeviceProfile().mBadgeRenderer;
-            PopupContainerWithArrow popup = PopupContainerWithArrow.getOpen(mLauncher);
-            if (popup != null) {
-                popup.updateNotificationHeader(badgeInfo, itemInfo);
-            }
-
             boolean wasBadged = mBadgeInfo != null;
-            boolean isBadged = badgeInfo != null;
+            mBadgeInfo = mLauncher.getPopupDataProvider().getBadgeInfoForItem(itemInfo);
+            boolean isBadged = mBadgeInfo != null;
             float newBadgeScale = isBadged ? 1f : 0;
-            mBadgeInfo = badgeInfo;
-            mBadgeRenderer = badgeRenderer;
+            mBadgeRenderer = mLauncher.getDeviceProfile().mBadgeRenderer;
             if (wasBadged || isBadged) {
                 mIconPalette = ((FastBitmapDrawable) mIcon).getIconPalette();
                 // Animate when a badge is first added or when it is removed.
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index dbf535a..3c29f5e 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1576,6 +1576,11 @@
             public void run() {
                 mWorkspace.updateIconBadges(updatedBadges);
                 mAppsView.updateIconBadges(updatedBadges);
+
+                PopupContainerWithArrow popup = PopupContainerWithArrow.getOpen(Launcher.this);
+                if (popup != null) {
+                    popup.updateNotificationHeader(updatedBadges);
+                }
             }
         };
         if (!waitUntilResume(r)) {
diff --git a/src/com/android/launcher3/popup/PopupContainerWithArrow.java b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
index ee64ef9..b00eb1f 100644
--- a/src/com/android/launcher3/popup/PopupContainerWithArrow.java
+++ b/src/com/android/launcher3/popup/PopupContainerWithArrow.java
@@ -77,6 +77,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static com.android.launcher3.userevent.nano.LauncherLogProto.ContainerType;
 import static com.android.launcher3.userevent.nano.LauncherLogProto.ItemType;
@@ -172,6 +173,8 @@
         final int arrowVerticalOffset = resources.getDimensionPixelSize(
                 R.dimen.popup_arrow_vertical_offset);
 
+        mOriginalIcon = originalIcon;
+
         // Add dummy views first, and populate with real info when ready.
         PopupPopulator.Item[] itemsToPopulate = PopupPopulator
                 .getItemsToPopulate(shortcutIds, notificationKeys, systemShortcuts);
@@ -200,9 +203,7 @@
                 ? Collections.EMPTY_LIST
                 : mShortcutsItemView.getSystemShortcutViews(reverseOrder);
         if (mNotificationItemView != null) {
-            BadgeInfo badgeInfo = mLauncher.getPopupDataProvider()
-                    .getBadgeInfoForItem(originalItemInfo);
-            updateNotificationHeader(badgeInfo, originalIcon);
+            updateNotificationHeader();
         }
 
         // Add the arrow.
@@ -212,7 +213,6 @@
 
         animateOpen();
 
-        mOriginalIcon = originalIcon;
         mLauncher.getDragController().addDragListener(this);
         mOriginalIcon.forceHideBadge(true);
 
@@ -551,21 +551,22 @@
     }
 
     /**
-     * Updates the notification header to reflect the badge info. Since this can be called
-     * for any badge info (not necessarily the one associated with this app), we first
-     * check that the ItemInfo matches the one of this popup.
+     * Updates the notification header if the original icon's badge updated.
      */
-    public void updateNotificationHeader(BadgeInfo badgeInfo, ItemInfo originalItemInfo) {
-        if (originalItemInfo != mOriginalIcon.getTag()) {
-            return;
+    public void updateNotificationHeader(Set<PackageUserKey> updatedBadges) {
+        ItemInfo itemInfo = (ItemInfo) mOriginalIcon.getTag();
+        PackageUserKey packageUser = PackageUserKey.fromItemInfo(itemInfo);
+        if (updatedBadges.contains(packageUser)) {
+            updateNotificationHeader();
         }
-        updateNotificationHeader(badgeInfo, mOriginalIcon);
     }
 
-    private void updateNotificationHeader(BadgeInfo badgeInfo, BubbleTextView originalIcon) {
+    private void updateNotificationHeader() {
+        ItemInfo itemInfo = (ItemInfo) mOriginalIcon.getTag();
+        BadgeInfo badgeInfo = mLauncher.getPopupDataProvider().getBadgeInfoForItem(itemInfo);
         if (mNotificationItemView != null && badgeInfo != null) {
-            IconPalette palette = originalIcon.getIcon() instanceof FastBitmapDrawable
-                    ? ((FastBitmapDrawable) originalIcon.getIcon()).getIconPalette()
+            IconPalette palette = mOriginalIcon.getIcon() instanceof FastBitmapDrawable
+                    ? ((FastBitmapDrawable) mOriginalIcon.getIcon()).getIconPalette()
                     : null;
             mNotificationItemView.updateHeader(badgeInfo.getNotificationCount(), palette);
         }