perf(ui): coalesce persisted-store writes before they hit host storage (#3299)

Author: tatoalo

packages/ui/src/features/settings/settingsStore.test.ts

@@ -1,4 +1,7 @@
-import { registerRendererStateStorage } from "@posthog/ui/shell/rendererStorage";
+import {
+  flushRendererStateWrites,
+  registerRendererStateStorage,
+} from "@posthog/ui/shell/rendererStorage";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import {
   type CompletionSound,
@@ -12,6 +15,28 @@ const removeItem = vi.fn();
 
 registerRendererStateStorage({ getItem, setItem, removeItem });
 
+// Lands any coalesced write from the previous test on the old mocks (so a
+// pending value cannot leak into this test's reads or assertions), then
+// resets them.
+async function resetPersistenceMocks() {
+  await flushRendererStateWrites();
+  getItem.mockReset();
+  setItem.mockReset();
+  removeItem.mockReset();
+  getItem.mockResolvedValue(null);
+  setItem.mockResolvedValue(undefined);
+  removeItem.mockResolvedValue(undefined);
+}
+
+// Persisted writes are debounced; flush while polling so the assertion sees
+// the coalesced write as soon as the store has queued it.
+async function waitForPersistedWrite() {
+  await vi.waitFor(async () => {
+    await flushRendererStateWrites();
+    expect(setItem).toHaveBeenCalled();
+  });
+}
+
 // Runs before any test mutates the store singleton, so getState() still
 // reflects the initial values.
 describe("feature settingsStore defaults", () => {
@@ -25,13 +50,8 @@ describe("feature settingsStore defaults", () => {
 });
 
 describe("feature settingsStore cloud selections", () => {
-  beforeEach(() => {
-    getItem.mockReset();
-    setItem.mockReset();
-    removeItem.mockReset();
-    getItem.mockResolvedValue(null);
-    setItem.mockResolvedValue(undefined);
-    removeItem.mockResolvedValue(undefined);
+  beforeEach(async () => {
+    await resetPersistenceMocks();
 
     useSettingsStore.setState({
       allowBypassPermissions: false,
@@ -43,9 +63,7 @@ describe("feature settingsStore cloud selections", () => {
   it("persists the last used cloud repository", async () => {
     useSettingsStore.getState().setLastUsedCloudRepository("posthog/posthog");
 
-    await vi.waitFor(() => {
-      expect(setItem).toHaveBeenCalled();
-    });
+    await waitForPersistedWrite();
 
     const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1];
     const persisted = JSON.parse(lastCall[1]);
@@ -82,9 +100,7 @@ describe("feature settingsStore cloud selections", () => {
       },
     });
 
-    await vi.waitFor(() => {
-      expect(setItem).toHaveBeenCalled();
-    });
+    await waitForPersistedWrite();
 
     const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1];
     const persisted = JSON.parse(lastCall[1]);
@@ -170,13 +186,8 @@ describe("feature settingsStore cloud selections", () => {
 });
 
 describe("feature settingsStore custom sounds", () => {
-  beforeEach(() => {
-    getItem.mockReset();
-    setItem.mockReset();
-    removeItem.mockReset();
-    getItem.mockResolvedValue(null);
-    setItem.mockResolvedValue(undefined);
-    removeItem.mockResolvedValue(undefined);
+  beforeEach(async () => {
+    await resetPersistenceMocks();
 
     useSettingsStore.setState({ customSounds: [], completionSound: "none" });
   });
@@ -239,9 +250,7 @@ describe("feature settingsStore custom sounds", () => {
   it("persists custom sounds", async () => {
     useSettingsStore.getState().addCustomSound(sound);
 
-    await vi.waitFor(() => {
-      expect(setItem).toHaveBeenCalled();
-    });
+    await waitForPersistedWrite();
 
     const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1];
     const persisted = JSON.parse(lastCall[1]);
@@ -284,13 +293,8 @@ describe("feature settingsStore custom sounds", () => {
 });
 
 describe("feature settingsStore terminal font", () => {
-  beforeEach(() => {
-    getItem.mockReset();
-    setItem.mockReset();
-    removeItem.mockReset();
-    getItem.mockResolvedValue(null);
-    setItem.mockResolvedValue(undefined);
-    removeItem.mockResolvedValue(undefined);
+  beforeEach(async () => {
+    await resetPersistenceMocks();
 
     useSettingsStore.setState({
       terminalFont: "berkeley-mono",
@@ -307,9 +311,7 @@ describe("feature settingsStore terminal font", () => {
     useSettingsStore.getState().setTerminalFont("custom");
     useSettingsStore.getState().setTerminalCustomFontFamily("Fira Code");
 
-    await vi.waitFor(() => {
-      expect(setItem).toHaveBeenCalled();
-    });
+    await waitForPersistedWrite();
 
     const lastCall = setItem.mock.calls[setItem.mock.calls.length - 1];
     const persisted = JSON.parse(lastCall[1]);

packages/ui/src/shell/rendererStorage.test.ts

@@ -1,6 +1,7 @@
-import { describe, expect, it, vi } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
+import { WRITE_DEBOUNCE_MS } from "./rendererStorage";
 
 type RendererStorageModule = typeof import("./rendererStorage");
 
@@ -9,6 +10,15 @@ async function importFreshRendererStorage(): Promise<RendererStorageModule> {
   return await import("./rendererStorage");
 }
 
+/** Fresh module with fake timers and a backend already registered. */
+async function setupRegisteredBackend(data: Record<string, string> = {}) {
+  vi.useFakeTimers();
+  const module = await importFreshRendererStorage();
+  const backend = fakeBackend(data);
+  module.registerRendererStateStorage(backend);
+  return { module, storage: jsonStorageOf(module), backend };
+}
+
 function fakeBackend(data: Record<string, string>) {
   return {
     getItem: vi.fn(async (key: string) => data[key] ?? null),
@@ -30,6 +40,10 @@ function jsonStorageOf(module: RendererStorageModule) {
 }
 
 describe("rendererStorage", () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
   it("serves reads issued before the host registers its storage", async () => {
     const module = await importFreshRendererStorage();
     const storage = jsonStorageOf(module);
@@ -51,6 +65,7 @@ describe("rendererStorage", () => {
   });
 
   it("drops writes racing the initial read, then writes through", async () => {
+    vi.useFakeTimers();
     const module = await importFreshRendererStorage();
     const storage = jsonStorageOf(module);
     const backend = fakeBackend({
@@ -65,31 +80,35 @@ describe("rendererStorage", () => {
 
     module.registerRendererStateStorage(backend);
     await Promise.all([read, racingWrite]);
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
     expect(backend.setItem).not.toHaveBeenCalled();
 
     await storage.setItem("settings-storage", {
       state: { mode: "changed" },
       version: 0,
     });
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
     expect(backend.setItem).toHaveBeenCalledTimes(1);
   });
 
   it("passes writes through for keys that were never read", async () => {
+    vi.useFakeTimers();
     const module = await importFreshRendererStorage();
     const storage = jsonStorageOf(module);
     const backend = fakeBackend({});
 
-    const write = storage.setItem("write-only", {
+    await storage.setItem("write-only", {
       state: { value: 1 },
       version: 0,
     });
     module.registerRendererStateStorage(backend);
-    await write;
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
 
     expect(backend.setItem).toHaveBeenCalledTimes(1);
   });
 
   it("settles concurrent initial reads of the same key once", async () => {
+    vi.useFakeTimers();
     const module = await importFreshRendererStorage();
     const storage = jsonStorageOf(module);
     const backend = fakeBackend({
@@ -110,10 +129,12 @@ describe("rendererStorage", () => {
       state: { mode: "changed" },
       version: 0,
     });
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
     expect(backend.setItem).toHaveBeenCalledTimes(1);
   });
 
   it("marks a key settled when the initial read rejects so later writes pass", async () => {
+    vi.useFakeTimers();
     const module = await importFreshRendererStorage();
     const storage = jsonStorageOf(module);
     const backend = fakeBackend({});
@@ -127,6 +148,7 @@ describe("rendererStorage", () => {
       state: { mode: "changed" },
       version: 0,
     });
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
     expect(backend.setItem).toHaveBeenCalledTimes(1);
   });
 
@@ -196,12 +218,84 @@ describe("rendererStorage", () => {
     });
 
     useStore.setState({ defaultInitialTaskMode: "plan" });
-    await vi.waitFor(() => {
+    await vi.waitFor(async () => {
+      await module.flushRendererStateWrites();
       expect(backend.setItem).toHaveBeenCalled();
     });
     const persisted = JSON.parse(
       backend.setItem.mock.calls[backend.setItem.mock.calls.length - 1][1],
     );
     expect(persisted.state.defaultInitialTaskMode).toBe("plan");
   });
+
+  it("coalesces a burst of writes into one backend write with the latest value", async () => {
+    const { storage, backend } = await setupRegisteredBackend();
+
+    await storage.setItem("drafts", { state: { value: 1 }, version: 0 });
+    await storage.setItem("drafts", { state: { value: 2 }, version: 0 });
+    await storage.setItem("drafts", { state: { value: 3 }, version: 0 });
+
+    expect(backend.setItem).not.toHaveBeenCalled();
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
+    expect(backend.setItem).toHaveBeenCalledTimes(1);
+    const persisted = JSON.parse(backend.setItem.mock.calls[0][1]);
+    expect(persisted.state.value).toBe(3);
+  });
+
+  it("flushes at the max-wait bound during sustained writes", async () => {
+    const { storage, backend } = await setupRegisteredBackend();
+
+    // Write every 500ms for 6s: a plain trailing debounce would never fire,
+    // the max-wait bound forces a flush mid-burst.
+    for (let i = 0; i < 12; i++) {
+      await storage.setItem("drafts", { state: { value: i }, version: 0 });
+      await vi.advanceTimersByTimeAsync(500);
+    }
+    expect(backend.setItem.mock.calls.length).toBeGreaterThanOrEqual(1);
+
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
+    const last = backend.setItem.mock.calls.at(-1);
+    expect(JSON.parse(last?.[1] ?? "{}").state.value).toBe(11);
+  });
+
+  it("lands the pending coalesced write before serving a read", async () => {
+    const { storage, backend } = await setupRegisteredBackend({
+      drafts: JSON.stringify({ state: { value: "stale" }, version: 0 }),
+    });
+
+    await storage.setItem("drafts", { state: { value: "fresh" }, version: 0 });
+
+    await expect(storage.getItem("drafts")).resolves.toMatchObject({
+      state: { value: "fresh" },
+    });
+    expect(backend.setItem).toHaveBeenCalledTimes(1);
+
+    // The read consumed the pending write; the debounce timer must not fire
+    // a duplicate.
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
+    expect(backend.setItem).toHaveBeenCalledTimes(1);
+  });
+
+  it("cancels a pending write when the key is removed", async () => {
+    const { storage, backend } = await setupRegisteredBackend();
+
+    await storage.setItem("drafts", { state: { value: 1 }, version: 0 });
+    await storage.removeItem("drafts");
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
+
+    expect(backend.setItem).not.toHaveBeenCalled();
+    expect(backend.removeItem).toHaveBeenCalledTimes(1);
+  });
+
+  it("flushRendererStateWrites persists pending state immediately", async () => {
+    const { module, storage, backend } = await setupRegisteredBackend();
+
+    await storage.setItem("drafts", { state: { value: 1 }, version: 0 });
+    await module.flushRendererStateWrites();
+    expect(backend.setItem).toHaveBeenCalledTimes(1);
+
+    // The debounce timer was cancelled by the flush; no duplicate write.
+    await vi.advanceTimersByTimeAsync(WRITE_DEBOUNCE_MS);
+    expect(backend.setItem).toHaveBeenCalledTimes(1);
+  });
 });