From 1a57ad3fc5cb6ec93fc6b354367e6a0c1122e248 Mon Sep 17 00:00:00 2001 From: Andreas Coroiu Date: Tue, 19 Aug 2025 18:05:06 +0200 Subject: [PATCH] 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 --- .../src/platform/badge/badge-browser-api.ts | 35 ++++- .../src/platform/badge/badge.service.spec.ts | 138 +++++++++++++++++- .../src/platform/badge/badge.service.ts | 121 +++++++++------ .../badge/test/mock-badge-browser-api.ts | 24 ++- .../src/platform/browser/browser-api.ts | 9 ++ 5 files changed, 266 insertions(+), 61 deletions(-) diff --git a/apps/browser/src/platform/badge/badge-browser-api.ts b/apps/browser/src/platform/badge/badge-browser-api.ts index 097c610974..e8edcd23da 100644 --- 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; + // activeTabs$: Observable; setState(state: RawBadgeState, tabId?: number): Promise; getTabs(): Promise; + getActiveTabs(): Promise; } 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 { + return BrowserApi.getActiveTabs(); + } + constructor(private platformUtilsService: PlatformUtilsService) {} async setState(state: RawBadgeState, tabId?: number): Promise { diff --git a/apps/browser/src/platform/badge/badge.service.spec.ts b/apps/browser/src/platform/badge/badge.service.spec.ts index 52be2afa71..71b94e3a0e 100644 --- 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, + }); + }); + }); }); }); diff --git a/apps/browser/src/platform/badge/badge.service.ts b/apps/browser/src/platform/badge/badge.service.ts index b3831530e8..8283637084 100644 --- 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>; + private serviceState: GlobalState>; 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, 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 | 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); + } + } + } } /** diff --git a/apps/browser/src/platform/badge/test/mock-badge-browser-api.ts b/apps/browser/src/platform/badge/test/mock-badge-browser-api.ts index 4f91420b27..6708ae52ea 100644 --- 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 = {}; generalState?: RawBadgeState; tabs: number[] = []; + activeTabs: number[] = []; - setActiveTab(tabId: number) { + getActiveTabs(): Promise { + 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 { + setState = jest.fn().mockImplementation((state: RawBadgeState, tabId?: number): Promise => { if (tabId !== undefined) { this.specificStates[tabId] = state; } else { @@ -25,7 +43,7 @@ export class MockBadgeBrowserApi implements BadgeBrowserApi { } return Promise.resolve(); - } + }); getTabs(): Promise { return Promise.resolve(this.tabs); diff --git a/apps/browser/src/platform/browser/browser-api.ts b/apps/browser/src/platform/browser/browser-api.ts index d0bdaa504b..60a42078cf 100644 --- 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 { + return new Promise((resolve) => chrome.windows.getAll({ populate: true }, resolve)); + } + /** * Gets the current window or the window with the given id. *