fix(badge): [PM-24661] Improve performance of badge state calculation for large number of tabs

* refactor: rewrite the badge service to only calculate states for active tab

This also fixes an issue where the `difference` function didn't work and caused
all tabs to update for every single state update.

* fix: compilation issue

* feat: add error logging

* fix: linting

* fix: badge clearing on reload on firefox

* feat: optimize observable

* feat(wip): update all active tabs

tests are broken

* fix: existing tests

* feat: add new tests

(cherry picked from commit 1a57ad3)

Author: coroiu

apps/browser/src/platform/badge/badge-browser-api.ts

@@ -1,4 +1,4 @@
-import { map, Observable } from "rxjs";
+import { concat, defer, filter, map, merge, Observable, shareReplay, switchMap } from "rxjs";
 
 import { PlatformUtilsService } from "@bitwarden/common/platform/abstractions/platform-utils.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
@@ -17,19 +17,48 @@ export interface RawBadgeState {
 
 export interface BadgeBrowserApi {
   activeTab$: Observable<chrome.tabs.TabActiveInfo | undefined>;
+  // activeTabs$: Observable<chrome.tabs.Tab[]>;
 
   setState(state: RawBadgeState, tabId?: number): Promise<void>;
   getTabs(): Promise<number[]>;
+  getActiveTabs(): Promise<chrome.tabs.Tab[]>;
 }
 
 export class DefaultBadgeBrowserApi implements BadgeBrowserApi {
   private badgeAction = BrowserApi.getBrowserAction();
   private sidebarAction = BrowserApi.getSidebarAction(self);
 
-  activeTab$ = fromChromeEvent(chrome.tabs.onActivated).pipe(
-    map(([tabActiveInfo]) => tabActiveInfo),
+  private onTabActivated$ = fromChromeEvent(chrome.tabs.onActivated).pipe(
+    switchMap(async ([activeInfo]) => activeInfo),
+    shareReplay({ bufferSize: 1, refCount: true }),
   );
 
+  activeTab$ = concat(
+    defer(async () => {
+      const currentTab = await BrowserApi.getTabFromCurrentWindow();
+      if (currentTab == null || currentTab.id === undefined) {
+        return undefined;
+      }
+
+      return { tabId: currentTab.id, windowId: currentTab.windowId };
+    }),
+    merge(
+      this.onTabActivated$,
+      fromChromeEvent(chrome.tabs.onUpdated).pipe(
+        filter(
+          ([_, changeInfo]) =>
+            // Only emit if the url was updated
+            changeInfo.url != undefined,
+        ),
+        map(([tabId, _changeInfo, tab]) => ({ tabId, windowId: tab.windowId })),
+      ),
+    ),
+  ).pipe(shareReplay({ bufferSize: 1, refCount: true }));
+
+  getActiveTabs(): Promise<chrome.tabs.Tab[]> {
+    return BrowserApi.getActiveTabs();
+  }
+
   constructor(private platformUtilsService: PlatformUtilsService) {}
 
   async setState(state: RawBadgeState, tabId?: number): Promise<void> {

apps/browser/src/platform/badge/badge.service.spec.ts

@@ -37,8 +37,9 @@ describe("BadgeService", () => {
 
     describe("given a single tab is open", () => {
       beforeEach(() => {
-        badgeApi.tabs = [1];
-        badgeApi.setActiveTab(tabId);
+        badgeApi.tabs = [tabId];
+        badgeApi.setActiveTabs([tabId]);
+        badgeApi.setLastActivatedTab(tabId);
         badgeServiceSubscription = badgeService.startListening();
       });
 
@@ -187,17 +188,18 @@ describe("BadgeService", () => {
       });
     });
 
-    describe("given multiple tabs are open", () => {
+    describe("given multiple tabs are open, only one active", () => {
       const tabId = 1;
       const tabIds = [1, 2, 3];
 
       beforeEach(() => {
         badgeApi.tabs = tabIds;
-        badgeApi.setActiveTab(tabId);
+        badgeApi.setActiveTabs([tabId]);
+        badgeApi.setLastActivatedTab(tabId);
         badgeServiceSubscription = badgeService.startListening();
       });
 
-      it("sets state for each tab when no other state has been set", async () => {
+      it("sets general state for active tab when no other state has been set", async () => {
         const state: BadgeState = {
           text: "text",
           backgroundColor: "color",
@@ -213,6 +215,67 @@ describe("BadgeService", () => {
           3: undefined,
         });
       });
+
+      it("only updates the active tab when setting state", async () => {
+        const state: BadgeState = {
+          text: "text",
+          backgroundColor: "color",
+          icon: BadgeIcon.Locked,
+        };
+        badgeApi.setState.mockReset();
+
+        await badgeService.setState("state-1", BadgeStatePriority.Default, state, tabId);
+        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        expect(badgeApi.setState).toHaveBeenCalledTimes(1);
+      });
+    });
+
+    describe("given multiple tabs are open and multiple are active", () => {
+      const activeTabIds = [1, 2];
+      const tabIds = [1, 2, 3];
+
+      beforeEach(() => {
+        badgeApi.tabs = tabIds;
+        badgeApi.setActiveTabs(activeTabIds);
+        badgeApi.setLastActivatedTab(1);
+        badgeServiceSubscription = badgeService.startListening();
+      });
+
+      it("sets general state for active tabs when no other state has been set", async () => {
+        const state: BadgeState = {
+          text: "text",
+          backgroundColor: "color",
+          icon: BadgeIcon.Locked,
+        };
+
+        await badgeService.setState("state-name", BadgeStatePriority.Default, state);
+
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        expect(badgeApi.specificStates).toEqual({
+          1: state,
+          2: state,
+          3: undefined,
+        });
+      });
+
+      it("only updates the active tabs when setting general state", async () => {
+        const state: BadgeState = {
+          text: "text",
+          backgroundColor: "color",
+          icon: BadgeIcon.Locked,
+        };
+        badgeApi.setState.mockReset();
+
+        await badgeService.setState("state-1", BadgeStatePriority.Default, state, 1);
+        await badgeService.setState("state-2", BadgeStatePriority.Default, state, 2);
+        await badgeService.setState("state-3", BadgeStatePriority.Default, state, 3);
+
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        expect(badgeApi.setState).toHaveBeenCalledTimes(2);
+      });
     });
   });
 
@@ -222,7 +285,8 @@ describe("BadgeService", () => {
 
       beforeEach(() => {
         badgeApi.tabs = [tabId];
-        badgeApi.setActiveTab(tabId);
+        badgeApi.setActiveTabs([tabId]);
+        badgeApi.setLastActivatedTab(tabId);
         badgeServiceSubscription = badgeService.startListening();
       });
 
@@ -491,13 +555,14 @@ describe("BadgeService", () => {
       });
     });
 
-    describe("given multiple tabs are open", () => {
+    describe("given multiple tabs are open, only one active", () => {
       const tabId = 1;
       const tabIds = [1, 2, 3];
 
       beforeEach(() => {
         badgeApi.tabs = tabIds;
-        badgeApi.setActiveTab(tabId);
+        badgeApi.setActiveTabs([tabId]);
+        badgeApi.setLastActivatedTab(tabId);
         badgeServiceSubscription = badgeService.startListening();
       });
 
@@ -528,5 +593,62 @@ describe("BadgeService", () => {
         });
       });
     });
+
+    describe("given multiple tabs are open and multiple are active", () => {
+      const tabId = 1;
+      const activeTabIds = [1, 2];
+      const tabIds = [1, 2, 3];
+
+      beforeEach(() => {
+        badgeApi.tabs = tabIds;
+        badgeApi.setActiveTabs(activeTabIds);
+        badgeApi.setLastActivatedTab(tabId);
+        badgeServiceSubscription = badgeService.startListening();
+      });
+
+      it("sets general state for all active tabs when no other state has been set", async () => {
+        const generalState: BadgeState = {
+          text: "general-text",
+          backgroundColor: "general-color",
+          icon: BadgeIcon.Unlocked,
+        };
+
+        await badgeService.setState("general-state", BadgeStatePriority.Default, generalState);
+
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        expect(badgeApi.specificStates).toEqual({
+          [tabIds[0]]: generalState,
+          [tabIds[1]]: generalState,
+          [tabIds[2]]: undefined,
+        });
+      });
+
+      it("sets tab-specific state for provided tab", async () => {
+        const generalState: BadgeState = {
+          text: "general-text",
+          backgroundColor: "general-color",
+          icon: BadgeIcon.Unlocked,
+        };
+        const specificState: BadgeState = {
+          text: "tab-text",
+          icon: BadgeIcon.Locked,
+        };
+
+        await badgeService.setState("general-state", BadgeStatePriority.Default, generalState);
+        await badgeService.setState(
+          "tab-state",
+          BadgeStatePriority.Default,
+          specificState,
+          tabIds[0],
+        );
+
+        await new Promise((resolve) => setTimeout(resolve, 0));
+        expect(badgeApi.specificStates).toEqual({
+          [tabIds[0]]: { ...specificState, backgroundColor: "general-color" },
+          [tabIds[1]]: generalState,
+          [tabIds[2]]: undefined,
+        });
+      });
+    });
   });
 });