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
2025-12-06 00:13:28 +00:00 · 2025-08-19 18:05:06 +02:00
parent 07421da8c5
commit 1a57ad3fc5
5 changed files with 266 additions and 61 deletions
--- a/apps/browser/src/platform/badge/badge-browser-api.ts
+++ b/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> {
--- a/apps/browser/src/platform/badge/badge.service.spec.ts
+++ b/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,
+        });
+      });
+    });
  });
 });
--- a/apps/browser/src/platform/badge/badge.service.ts
+++ b/apps/browser/src/platform/badge/badge.service.ts
@@ -1,13 +1,4 @@
-import {
-  combineLatest,
-  concatMap,
-  distinctUntilChanged,
-  filter,
-  map,
-  pairwise,
-  startWith,
-  Subscription,
-} from "rxjs";
+import { concatMap, filter, Subscription, withLatestFrom } 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({
-      states: this.states.state$.pipe(
-        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)),
-    })
+    // React to tab changes
+    return this.badgeApi.activeTab$
      .pipe(
-        concatMap(async ({ states, activeTab }) => {
-          const changed = [...states.removed, ...states.added];
-
-          // If the active tab wasn't changed, we don't need to update the badge.
-          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);
-          }
+        withLatestFrom(this.serviceState.state$),
+        filter(([activeTab]) => activeTab != undefined),
+        concatMap(async ([activeTab, serviceState]) => {
+          await this.updateBadge(serviceState, activeTab!.tabId);
        }),
      )
      .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);
+      }
+    }
+  }
 }

 /**
--- a/apps/browser/src/platform/badge/test/mock-badge-browser-api.ts
+++ b/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);
--- a/apps/browser/src/platform/browser/browser-api.ts
+++ b/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.
   *