Ps/run-foreground-derived-state-in-zone (#7861)

* Sync derived state through memory storage

* Run foreground derived state in NgZone

* fix tests

apps/browser/src/platform/state/background-derived-state.ts

@@ -1,11 +1,10 @@
-import { Observable, Subject, Subscription } from "rxjs";
+import { Observable, Subscription } from "rxjs";
 import { Jsonify } from "type-fest";
 
 import {
   AbstractStorageService,
   ObservableStorageService,
 } from "@bitwarden/common/platform/abstractions/storage.service";
-import { Utils } from "@bitwarden/common/platform/misc/utils";
 import { DeriveDefinition } from "@bitwarden/common/platform/state";
 // eslint-disable-next-line import/no-restricted-paths -- extending this class for this client
 import { DefaultDerivedState } from "@bitwarden/common/platform/state/implementations/default-derived-state";
@@ -18,10 +17,7 @@ export class BackgroundDerivedState<
   TTo,
   TDeps extends DerivedStateDependencies,
 > extends DefaultDerivedState<TFrom, TTo, TDeps> {
-  private portSubscriptions: Map<
-    chrome.runtime.Port,
-    { subscription: Subscription; delaySubject: Subject<void> }
-  > = new Map();
+  private portSubscriptions: Map<chrome.runtime.Port, Subscription> = new Map();
 
   constructor(
     parentState$: Observable<TFrom>,
@@ -40,34 +36,15 @@ export class BackgroundDerivedState<
 
       const listenerCallback = this.onMessageFromForeground.bind(this);
       port.onDisconnect.addListener(() => {
-        const { subscription, delaySubject } = this.portSubscriptions.get(port) ?? {
-          subscription: null,
-          delaySubject: null,
-        };
-        subscription?.unsubscribe();
-        delaySubject?.complete();
+        this.portSubscriptions.get(port)?.unsubscribe();
         this.portSubscriptions.delete(port);
         port.onMessage.removeListener(listenerCallback);
       });
       port.onMessage.addListener(listenerCallback);
 
-      const delaySubject = new Subject<void>();
-      const stateSubscription = this.state$.subscribe((state) => {
-        // delay to allow the foreground to connect. This may just be needed for testing
-        setTimeout(() => {
-          // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
-          // eslint-disable-next-line @typescript-eslint/no-floating-promises
-          this.sendNewMessage(
-            {
-              action: "nextState",
-              data: JSON.stringify(state),
-            },
-            port,
-          );
-        }, 0);
-      });
+      const stateSubscription = this.state$.subscribe();
 
-      this.portSubscriptions.set(port, { subscription: stateSubscription, delaySubject });
+      this.portSubscriptions.set(port, stateSubscription);
     });
   }
 
@@ -93,22 +70,6 @@ export class BackgroundDerivedState<
     }
   }
 
-  private async sendNewMessage(
-    message: Omit<DerivedStateMessage, "originator" | "id">,
-    port: chrome.runtime.Port,
-  ) {
-    const id = Utils.newGuid();
-    // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
-    // eslint-disable-next-line @typescript-eslint/no-floating-promises
-    this.sendMessage(
-      {
-        ...message,
-        id: id,
-      },
-      port,
-    );
-  }
-
   private async sendResponse(
     originalMessage: DerivedStateMessage,
     response: Omit<DerivedStateMessage, "originator" | "id">,

apps/browser/src/platform/state/derived-state-interactions.spec.ts

@@ -3,8 +3,10 @@
  * @jest-environment ../../libs/shared/test.environment.ts
  */
 
+import { NgZone } from "@angular/core";
 import { FakeStorageService } from "@bitwarden/common/../spec/fake-storage.service";
 import { awaitAsync, trackEmissions } from "@bitwarden/common/../spec/utils";
+import { mock } from "jest-mock-extended";
 import { Subject, firstValueFrom } from "rxjs";
 
 import { DeriveDefinition } from "@bitwarden/common/platform/state";
@@ -22,12 +24,20 @@ const deriveDefinition = new DeriveDefinition(stateDefinition, "test", {
   deserializer: (dateString: string) => (dateString == null ? null : new Date(dateString)),
 });
 
+// Mock out the runInsideAngular operator so we don't have to deal with zone.js
+jest.mock("../browser/run-inside-angular.operator", () => {
+  return {
+    runInsideAngular: (ngZone: any) => (source: any) => source,
+  };
+});
+
 describe("foreground background derived state interactions", () => {
   let foreground: ForegroundDerivedState<Date>;
   let background: BackgroundDerivedState<string, Date, Record<string, unknown>>;
   let parentState$: Subject<string>;
   let memoryStorage: FakeStorageService;
   const initialParent = "2020-01-01";
+  const ngZone = mock<NgZone>();
 
   beforeEach(() => {
     mockPorts();
@@ -35,7 +45,7 @@ describe("foreground background derived state interactions", () => {
     memoryStorage = new FakeStorageService();
 
     background = new BackgroundDerivedState(parentState$, deriveDefinition, memoryStorage, {});
-    foreground = new ForegroundDerivedState(deriveDefinition);
+    foreground = new ForegroundDerivedState(deriveDefinition, memoryStorage, ngZone);
   });
 
   afterEach(() => {
@@ -50,12 +60,12 @@ describe("foreground background derived state interactions", () => {
     parentState$.next(initialParent);
     await awaitAsync(10);
 
-    expect(foregroundEmissions).toEqual([new Date(initialParent)]);
     expect(backgroundEmissions).toEqual([new Date(initialParent)]);
+    expect(foregroundEmissions).toEqual([new Date(initialParent)]);
   });
 
   it("should initialize a late-connected foreground", async () => {
-    const newForeground = new ForegroundDerivedState(deriveDefinition);
+    const newForeground = new ForegroundDerivedState(deriveDefinition, memoryStorage, ngZone);
     const backgroundEmissions = trackEmissions(background.state$);
     parentState$.next(initialParent);
     await awaitAsync();
@@ -74,7 +84,7 @@ describe("foreground background derived state interactions", () => {
 
       // FIXME: Verify that this floating promise is intentional. If it is, add an explanatory comment and ensure there is proper error handling.
       // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      foreground.forceValue(new Date(dateString));
+      await foreground.forceValue(new Date(dateString));
       await awaitAsync();
 
       expect(emissions).toEqual([new Date(dateString)]);

apps/browser/src/platform/state/foreground-derived-state.provider.ts

@@ -1,5 +1,10 @@
+import { NgZone } from "@angular/core";
 import { Observable } from "rxjs";
 
+import {
+  AbstractStorageService,
+  ObservableStorageService,
+} from "@bitwarden/common/platform/abstractions/storage.service";
 import { DeriveDefinition, DerivedState } from "@bitwarden/common/platform/state";
 // eslint-disable-next-line import/no-restricted-paths -- extending this class for this client
 import { DefaultDerivedStateProvider } from "@bitwarden/common/platform/state/implementations/default-derived-state.provider";
@@ -8,11 +13,17 @@ import { DerivedStateDependencies } from "@bitwarden/common/src/types/state";
 import { ForegroundDerivedState } from "./foreground-derived-state";
 
 export class ForegroundDerivedStateProvider extends DefaultDerivedStateProvider {
+  constructor(
+    memoryStorage: AbstractStorageService & ObservableStorageService,
+    private ngZone: NgZone,
+  ) {
+    super(memoryStorage);
+  }
   override buildDerivedState<TFrom, TTo, TDeps extends DerivedStateDependencies>(
     _parentState$: Observable<TFrom>,
     deriveDefinition: DeriveDefinition<TFrom, TTo, TDeps>,
     _dependencies: TDeps,
   ): DerivedState<TTo> {
-    return new ForegroundDerivedState(deriveDefinition);
+    return new ForegroundDerivedState(deriveDefinition, this.memoryStorage, this.ngZone);
   }
 }

apps/browser/src/platform/state/foreground-derived-state.spec.ts

@@ -1,4 +1,12 @@
-import { awaitAsync } from "@bitwarden/common/../spec/utils";
+/**
+ * need to update test environment so structuredClone works appropriately
+ * @jest-environment ../../libs/shared/test.environment.ts
+ */
+
+import { NgZone } from "@angular/core";
+import { awaitAsync, trackEmissions } from "@bitwarden/common/../spec";
+import { FakeStorageService } from "@bitwarden/common/../spec/fake-storage.service";
+import { mock } from "jest-mock-extended";
 
 import { DeriveDefinition } from "@bitwarden/common/platform/state";
 // eslint-disable-next-line import/no-restricted-paths -- needed to define a derive definition
@@ -15,12 +23,23 @@ const deriveDefinition = new DeriveDefinition(stateDefinition, "test", {
   cleanupDelayMs: 1,
 });
 
+// Mock out the runInsideAngular operator so we don't have to deal with zone.js
+jest.mock("../browser/run-inside-angular.operator", () => {
+  return {
+    runInsideAngular: (ngZone: any) => (source: any) => source,
+  };
+});
+
 describe("ForegroundDerivedState", () => {
   let sut: ForegroundDerivedState<Date>;
+  let memoryStorage: FakeStorageService;
+  const ngZone = mock<NgZone>();
 
   beforeEach(() => {
+    memoryStorage = new FakeStorageService();
+    memoryStorage.internalUpdateValuesRequireDeserialization(true);
     mockPorts();
-    sut = new ForegroundDerivedState(deriveDefinition);
+    sut = new ForegroundDerivedState(deriveDefinition, memoryStorage, ngZone);
   });
 
   afterEach(() => {
@@ -48,14 +67,17 @@ describe("ForegroundDerivedState", () => {
     expect(sut["port"]).toBeNull();
   });
 
-  it("should complete its replay subject when torn down", async () => {
-    const subscription = sut.state$.subscribe();
+  it("should emit when the memory storage updates", async () => {
+    const dateString = "2020-01-01";
+    const emissions = trackEmissions(sut.state$);
 
-    const completeSpy = jest.spyOn(sut["replaySubject"], "complete");
-    subscription.unsubscribe();
-    // wait for the cleanup delay
-    await awaitAsync(deriveDefinition.cleanupDelayMs * 2);
+    await memoryStorage.save(deriveDefinition.storageKey, {
+      derived: true,
+      value: new Date(dateString),
+    });
 
-    expect(completeSpy).toHaveBeenCalled();
+    await awaitAsync();
+
+    expect(emissions).toEqual([new Date(dateString)]);
   });
 });

apps/browser/src/platform/state/foreground-derived-state.ts

@@ -1,3 +1,4 @@
+import { NgZone } from "@angular/core";
 import {
   Observable,
   ReplaySubject,
@@ -5,36 +6,67 @@ import {
   filter,
   firstValueFrom,
   map,
+  merge,
+  of,
   share,
+  switchMap,
   tap,
   timer,
 } from "rxjs";
+import { Jsonify, JsonObject } from "type-fest";
 
+import {
+  AbstractStorageService,
+  ObservableStorageService,
+} from "@bitwarden/common/platform/abstractions/storage.service";
 import { Utils } from "@bitwarden/common/platform/misc/utils";
 import { DeriveDefinition, DerivedState } from "@bitwarden/common/platform/state";
 import { DerivedStateDependencies } from "@bitwarden/common/types/state";
 
 import { fromChromeEvent } from "../browser/from-chrome-event";
+import { runInsideAngular } from "../browser/run-inside-angular.operator";
 
 export class ForegroundDerivedState<TTo> implements DerivedState<TTo> {
+  private storageKey: string;
   private port: chrome.runtime.Port;
-  // For testing purposes
-  private replaySubject: ReplaySubject<TTo>;
   private backgroundResponses$: Observable<DerivedStateMessage>;
   state$: Observable<TTo>;
 
-  constructor(private deriveDefinition: DeriveDefinition<unknown, TTo, DerivedStateDependencies>) {
-    this.state$ = defer(() => this.initializePort()).pipe(
-      filter((message) => message.action === "nextState"),
-      map((message) => this.hydrateNext(message.data)),
-      share({
-        connector: () => {
-          this.replaySubject = new ReplaySubject<TTo>(1);
-          return this.replaySubject;
-        },
-        resetOnRefCountZero: () =>
-          timer(this.deriveDefinition.cleanupDelayMs).pipe(tap(() => this.tearDown())),
+  constructor(
+    private deriveDefinition: DeriveDefinition<unknown, TTo, DerivedStateDependencies>,
+    private memoryStorage: AbstractStorageService & ObservableStorageService,
+    private ngZone: NgZone,
+  ) {
+    this.storageKey = deriveDefinition.storageKey;
+
+    const initialStorageGet$ = defer(() => {
+      return this.getStoredValue();
+    }).pipe(
+      filter((s) => s.derived),
+      map((s) => s.value),
+    );
+
+    const latestStorage$ = this.memoryStorage.updates$.pipe(
+      filter((s) => s.key === this.storageKey),
+      switchMap(async (storageUpdate) => {
+        if (storageUpdate.updateType === "remove") {
+          return null;
+        }
+
+        return await this.getStoredValue();
       }),
+      filter((s) => s.derived),
+      map((s) => s.value),
+    );
+
+    this.state$ = defer(() => of(this.initializePort())).pipe(
+      switchMap(() => merge(initialStorageGet$, latestStorage$)),
+      share({
+        connector: () => new ReplaySubject<TTo>(1),
+        resetOnRefCountZero: () =>
+          timer(this.deriveDefinition.cleanupDelayMs).pipe(tap(() => this.tearDownPort())),
+      }),
+      runInsideAngular(this.ngZone),
     );
   }
 
@@ -51,7 +83,7 @@ export class ForegroundDerivedState<TTo> implements DerivedState<TTo> {
     return value;
   }
 
-  private initializePort(): Observable<DerivedStateMessage> {
+  private initializePort() {
     if (this.port != null) {
       return;
     }
@@ -88,11 +120,6 @@ export class ForegroundDerivedState<TTo> implements DerivedState<TTo> {
     });
   }
 
-  private hydrateNext(value: string): TTo {
-    const jsonObj = JSON.parse(value);
-    return this.deriveDefinition.deserialize(jsonObj);
-  }
-
   private tearDownPort() {
     if (this.port == null) {
       return;
@@ -103,8 +130,27 @@ export class ForegroundDerivedState<TTo> implements DerivedState<TTo> {
     this.backgroundResponses$ = null;
   }
 
-  private tearDown() {
-    this.tearDownPort();
-    this.replaySubject.complete();
+  protected async getStoredValue(): Promise<{ derived: boolean; value: TTo | null }> {
+    if (this.memoryStorage.valuesRequireDeserialization) {
+      const storedJson = await this.memoryStorage.get<
+        Jsonify<{ derived: true; value: JsonObject }>
+      >(this.storageKey);
+
+      if (!storedJson?.derived) {
+        return { derived: false, value: null };
+      }
+
+      const value = this.deriveDefinition.deserialize(storedJson.value as any);
+
+      return { derived: true, value };
+    } else {
+      const stored = await this.memoryStorage.get<{ derived: true; value: TTo }>(this.storageKey);
+
+      if (!stored?.derived) {
+        return { derived: false, value: null };
+      }
+
+      return { derived: true, value: stored.value };
+    }
   }
 }

apps/browser/src/popup/services/services.module.ts

@@ -1,4 +1,4 @@
-import { APP_INITIALIZER, LOCALE_ID, NgModule } from "@angular/core";
+import { APP_INITIALIZER, LOCALE_ID, NgModule, NgZone } from "@angular/core";
 
 import { UnauthGuard as BaseUnauthGuardService } from "@bitwarden/angular/auth/guards";
 import { ThemingService } from "@bitwarden/angular/platform/services/theming/theming.service";
@@ -554,7 +554,7 @@ function getBgService<T>(service: keyof MainBackground) {
     {
       provide: DerivedStateProvider,
       useClass: ForegroundDerivedStateProvider,
-      deps: [OBSERVABLE_MEMORY_STORAGE],
+      deps: [OBSERVABLE_MEMORY_STORAGE, NgZone],
     },
   ],
 })

libs/common/src/platform/state/implementations/default-derived-state.spec.ts

@@ -73,12 +73,12 @@ describe("DefaultDerivedState", () => {
     await awaitAsync();
 
     expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(
-      new Date(dateString),
+      derivedValue(new Date(dateString)),
     );
     const calls = memoryStorage.mock.save.mock.calls;
     expect(calls.length).toBe(1);
     expect(calls[0][0]).toBe(deriveDefinition.buildCacheKey());
-    expect(calls[0][1]).toEqual(new Date(dateString));
+    expect(calls[0][1]).toEqual(derivedValue(new Date(dateString)));
   });
 
   describe("forceValue", () => {
@@ -94,7 +94,9 @@ describe("DefaultDerivedState", () => {
 
       it("should store the forced value", async () => {
         await sut.forceValue(forced);
-        expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(forced);
+        expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(
+          derivedValue(forced),
+        );
       });
     });
 
@@ -107,7 +109,9 @@ describe("DefaultDerivedState", () => {
 
       it("should store the forced value", async () => {
         await sut.forceValue(forced);
-        expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(forced);
+        expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(
+          derivedValue(forced),
+        );
       });
 
       it("should force the value", async () => {
@@ -150,7 +154,7 @@ describe("DefaultDerivedState", () => {
       await awaitAsync();
 
       expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(
-        new Date(newDate),
+        derivedValue(new Date(newDate)),
       );
 
       subscription.unsubscribe();
@@ -168,7 +172,7 @@ describe("DefaultDerivedState", () => {
       await awaitAsync();
 
       expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(
-        new Date(newDate),
+        derivedValue(new Date(newDate)),
       );
 
       subscription.unsubscribe();
@@ -176,7 +180,7 @@ describe("DefaultDerivedState", () => {
       await awaitAsync(cleanupDelayMs * 2);
 
       expect(memoryStorage.internalStore[deriveDefinition.buildCacheKey()]).toEqual(
-        new Date(newDate),
+        derivedValue(new Date(newDate)),
       );
     });
 
@@ -256,3 +260,7 @@ describe("DefaultDerivedState", () => {
     });
   });
 });
+
+function derivedValue<T>(value: T) {
+  return { derived: true, value };
+}

libs/common/src/platform/state/implementations/default-derived-state.ts

@@ -34,7 +34,7 @@ export class DefaultDerivedState<TFrom, TTo, TDeps extends DerivedStateDependenc
           derivedStateOrPromise = await derivedStateOrPromise;
         }
         const derivedState = derivedStateOrPromise;
-        await this.memoryStorage.save(this.storageKey, derivedState);
+        await this.storeValue(derivedState);
         return derivedState;
       }),
     );
@@ -58,8 +58,12 @@ export class DefaultDerivedState<TFrom, TTo, TDeps extends DerivedStateDependenc
   }
 
   async forceValue(value: TTo) {
-    await this.memoryStorage.save(this.storageKey, value);
+    await this.storeValue(value);
     this.forcedValueSubject.next(value);
     return value;
   }
+
+  private storeValue(value: TTo) {
+    return this.memoryStorage.save(this.storageKey, { derived: true, value });
+  }
 }