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

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(
+  private onTabActivated$ = fromChromeEvent(chrome.tabs.onActivated).pipe(
-    map(([tabActiveInfo]) => tabActiveInfo),
+    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.tabs = [tabId];
-        badgeApi.setActiveTab(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,
+        });
+      });
+    });
   });
 });

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

@@ -1,13 +1,4 @@
-import {
+import { concatMap, filter, Subscription, withLatestFrom } from "rxjs";
-  combineLatest,
-  concatMap,
-  distinctUntilChanged,
-  filter,
-  map,
-  pairwise,
-  startWith,
-  Subscription,
-} from "rxjs";
 
 import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
 import {
@@ -17,7 +8,6 @@ import {
   StateProvider,
 } from "@bitwarden/common/platform/state";
 
-import { difference } from "./array-utils";
 import { BadgeBrowserApi, RawBadgeState } from "./badge-browser-api";
 import { DefaultBadgeState } from "./consts";
 import { BadgeStatePriority } from "./priority";
@@ -34,14 +24,14 @@ const BADGE_STATES = new KeyDefinition(BADGE_MEMORY, "badgeStates", {
 });
 
 export class BadgeService {
-  private states: GlobalState<Record<string, StateSetting>>;
+  private serviceState: GlobalState<Record<string, StateSetting>>;
 
   constructor(
     private stateProvider: StateProvider,
     private badgeApi: BadgeBrowserApi,
     private logService: LogService,
   ) {
-    this.states = this.stateProvider.getGlobal(BADGE_STATES);
+    this.serviceState = this.stateProvider.getGlobal(BADGE_STATES);
   }
 
   /**
@@ -49,44 +39,20 @@ export class BadgeService {
    * Without this the service will not be able to update the badge state.
    */
   startListening(): Subscription {
-    return combineLatest({
+    // React to tab changes
-      states: this.states.state$.pipe(
+    return this.badgeApi.activeTab$
-        startWith({}),
-        distinctUntilChanged(),
-        map((states) => new Set(states ? Object.values(states) : [])),
-        pairwise(),
-        map(([previous, current]) => {
-          const [removed, added] = difference(previous, current);
-          return { all: current, removed, added };
-        }),
-        filter(({ removed, added }) => removed.size > 0 || added.size > 0),
-      ),
-      activeTab: this.badgeApi.activeTab$.pipe(startWith(undefined)),
-    })
       .pipe(
-        concatMap(async ({ states, activeTab }) => {
+        withLatestFrom(this.serviceState.state$),
-          const changed = [...states.removed, ...states.added];
+        filter(([activeTab]) => activeTab != undefined),
-
+        concatMap(async ([activeTab, serviceState]) => {
-          // If the active tab wasn't changed, we don't need to update the badge.
+          await this.updateBadge(serviceState, activeTab!.tabId);
-          if (!changed.some((s) => s.tabId === activeTab?.tabId || s.tabId === undefined)) {
-            return;
-          }
-
-          try {
-            const state = this.calculateState(states.all, activeTab?.tabId);
-            await this.badgeApi.setState(state, activeTab?.tabId);
-          } catch (error) {
-            // This usually happens when the user opens a popout because of how the browser treats it
-            // as a tab in the same window but then won't let you set the badge state for it.
-            this.logService.warning("Failed to set badge state", error);
-          }
         }),
       )
       .subscribe({
-        error: (err: unknown) => {
+        error: (error: unknown) => {
           this.logService.error(
             "Fatal error in badge service observable, badge will fail to update",
-            err,
+            error,
           );
         },
       });
@@ -108,7 +74,12 @@ export class BadgeService {
    * @param tabId Limit this badge state to a specific tab. If this is not set, the state will be applied to all tabs.
    */
   async setState(name: string, priority: BadgeStatePriority, state: BadgeState, tabId?: number) {
-    await this.states.update((s) => ({ ...s, [name]: { priority, state, tabId } }));
+    const newServiceState = await this.serviceState.update((s) => ({
+      ...s,
+      [name]: { priority, state, tabId },
+    }));
+
+    await this.updateBadge(newServiceState, tabId);
   }
 
   /**
@@ -120,11 +91,21 @@ export class BadgeService {
    * @param name The name of the state to clear.
    */
   async clearState(name: string) {
-    await this.states.update((s) => {
+    let clearedState: StateSetting | undefined;
+
+    const newServiceState = await this.serviceState.update((s) => {
+      clearedState = s?.[name];
+
       const newStates = { ...s };
       delete newStates[name];
       return newStates;
     });
+
+    if (clearedState === undefined) {
+      return;
+    }
+    // const activeTabs = await firstValueFrom(this.badgeApi.activeTabs$);
+    await this.updateBadge(newServiceState, clearedState.tabId);
   }
 
   private calculateState(states: Set<StateSetting>, tabId?: number): RawBadgeState {
@@ -159,6 +140,52 @@ export class BadgeService {
       ...mergedState,
     };
   }
+
+  /**
+   * Common function deduplicating the logic for updating the badge with the current state.
+   * This will only update the badge if the active tab is the same as the tabId of the latest change.
+   * If the active tab is not set, it will not update the badge.
+   *
+   * @param activeTab The currently active tab.
+   * @param serviceState The current state of the badge service. If this is null or undefined, an empty set will be assumed.
+   * @param tabId Tab id for which the the latest state change applied to. Set this to activeTab.tabId to force an update.
+   */
+  private async updateBadge(
+    // activeTabs: chrome.tabs.Tab[],
+    serviceState: Record<string, StateSetting> | null | undefined,
+    tabId: number | undefined,
+  ) {
+    const activeTabs = await this.badgeApi.getActiveTabs();
+    if (tabId !== undefined && !activeTabs.some((tab) => tab.id === tabId)) {
+      return; // No need to update the badge if the state is not for the active tab.
+    }
+
+    const tabIdsToUpdate = tabId ? [tabId] : activeTabs.map((tab) => tab.id);
+
+    for (const tabId of tabIdsToUpdate) {
+      if (tabId === undefined) {
+        continue; // Skip if tab id is undefined.
+      }
+
+      const newBadgeState = this.calculateState(new Set(Object.values(serviceState ?? {})), tabId);
+      try {
+        await this.badgeApi.setState(newBadgeState, tabId);
+      } catch (error) {
+        this.logService.error("Failed to set badge state", error);
+      }
+    }
+
+    if (tabId === undefined) {
+      // If no tabId was provided we should also update the general badge state
+      const newBadgeState = this.calculateState(new Set(Object.values(serviceState ?? {})));
+
+      try {
+        await this.badgeApi.setState(newBadgeState, tabId);
+      } catch (error) {
+        this.logService.error("Failed to set general badge state", error);
+      }
+    }
+  }
 }
 
 /**

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

@@ -9,15 +9,33 @@ export class MockBadgeBrowserApi implements BadgeBrowserApi {
   specificStates: Record<number, RawBadgeState> = {};
   generalState?: RawBadgeState;
   tabs: number[] = [];
+  activeTabs: number[] = [];
 
-  setActiveTab(tabId: number) {
+  getActiveTabs(): Promise<chrome.tabs.Tab[]> {
+    return Promise.resolve(
+      this.activeTabs.map(
+        (tabId) =>
+          ({
+            id: tabId,
+            windowId: 1,
+            active: true,
+          }) as chrome.tabs.Tab,
+      ),
+    );
+  }
+
+  setActiveTabs(tabs: number[]) {
+    this.activeTabs = tabs;
+  }
+
+  setLastActivatedTab(tabId: number) {
     this._activeTab$.next({
       tabId,
       windowId: 1,
     });
   }
 
-  setState(state: RawBadgeState, tabId?: number): Promise<void> {
+  setState = jest.fn().mockImplementation((state: RawBadgeState, tabId?: number): Promise<void> => {
     if (tabId !== undefined) {
       this.specificStates[tabId] = state;
     } else {
@@ -25,7 +43,7 @@ export class MockBadgeBrowserApi implements BadgeBrowserApi {
     }
 
     return Promise.resolve();
-  }
+  });
 
   getTabs(): Promise<number[]> {
     return Promise.resolve(this.tabs);

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

@@ -32,6 +32,15 @@ export class BrowserApi {
     return BrowserApi.manifestVersion === expectedVersion;
   }
 
+  /**
+   * Gets all open browser windows, including their tabs.
+   *
+   * @returns A promise that resolves to an array of browser windows.
+   */
+  static async getWindows(): Promise<chrome.windows.Window[]> {
+    return new Promise((resolve) => chrome.windows.getAll({ populate: true }, resolve));
+  }
+
   /**
    * Gets the current window or the window with the given id.
    *