perf: Deduplicate concurrent Home Assistant registry list fetches (#2650)

src/ha/registry/device/index.ts

@@ -1,4 +1,5 @@
 import { errorToConsole } from '../../../utils/basic';
+import { OnceRunner } from '../../../utils/concurrency/once-runner';
 import type { HomeAssistant } from '../../types';
 import { homeAssistantWSRequest } from '../../ws-request';
 import {
@@ -10,7 +11,7 @@ import {
 
 export class DeviceRegistryManager {
   private _cache: DeviceCache;
-  private _fetchedDeviceList = false;
+  private _deviceListFetch = new OnceRunner();
 
   constructor(cache: DeviceCache) {
     this._cache = cache;
@@ -35,22 +36,21 @@ export class DeviceRegistryManager {
   }
 
   private async _fetchDeviceList(hass: HomeAssistant): Promise<void> {
-    if (this._fetchedDeviceList) {
-      return;
-    }
-
-    let deviceList: DeviceList | null = null;
     try {
-      deviceList = await homeAssistantWSRequest<DeviceList>(hass, deviceListSchema, {
+      await this._deviceListFetch.run(async () => {
+        const deviceList = await homeAssistantWSRequest<DeviceList>(
+          hass,
+          deviceListSchema,
+          {
             type: 'config/device_registry/list',
-      });
+          },
-    } catch (e) {
+        );
-      errorToConsole(e);
-      return;
-    }
         deviceList.forEach((device) => {
           this._cache.set(device.id, device);
         });
-    this._fetchedDeviceList = true;
+      });
+    } catch (e) {
+      errorToConsole(e);
+    }
   }
 }

src/ha/registry/entity/index.ts

@@ -1,4 +1,5 @@
 import { errorToConsole } from '../../../utils/basic.js';
+import { OnceRunner } from '../../../utils/concurrency/once-runner.js';
 import type { HomeAssistant } from '../../types.js';
 import { homeAssistantWSRequest } from '../../ws-request.js';
 import {
@@ -16,7 +17,7 @@ import {
 
 export class EntityRegistryManagerLive implements EntityRegistryManager {
   private _cache: EntityCache;
-  private _fetchedEntityList = false;
+  private _entityListFetch = new OnceRunner();
 
   constructor(cache: EntityCache) {
     this._cache = cache;
@@ -69,22 +70,21 @@ export class EntityRegistryManagerLive implements EntityRegistryManager {
   }
 
   public async fetchEntityList(hass: HomeAssistant): Promise<void> {
-    if (this._fetchedEntityList) {
-      return;
-    }
-
-    let entityList: EntityList | null = null;
     try {
-      entityList = await homeAssistantWSRequest<EntityList>(hass, entityListSchema, {
+      await this._entityListFetch.run(async () => {
+        const entityList = await homeAssistantWSRequest<EntityList>(
+          hass,
+          entityListSchema,
+          {
             type: 'config/entity_registry/list',
-      });
+          },
-    } catch (e) {
+        );
-      errorToConsole(e);
-      return;
-    }
         entityList.forEach((entity) => {
           this._cache.set(entity.entity_id, entity);
         });
-    this._fetchedEntityList = true;
+      });
+    } catch (e) {
+      errorToConsole(e);
+    }
   }
 }

src/utils/concurrency/once-runner.ts

@@ -0,0 +1,45 @@
+export type Work = () => Promise<void>;
+
+/**
+ * Runs asynchronous work at most once, and shares it while it is in flight.
+ *
+ * The first call starts the work. Callers that arrive while it is still running
+ * wait on that same run rather than starting their own, so ten concurrent
+ * callers make one request instead of ten identical ones. Once the work has
+ * succeeded, later calls return immediately without running it again.
+ *
+ * A failure is *not* remembered: every caller waiting on the failed run sees
+ * the rejection, and the next call starts a fresh attempt.
+ *
+ * Only the first caller's `work` ever runs. Callers that join an in-flight run
+ * have their own `work` discarded, so every caller must pass work that is
+ * interchangeable with the others'.
+ */
+export class OnceRunner {
+  private _succeeded = false;
+  private _inFlight: Promise<void> | null = null;
+
+  public async run(work: Work): Promise<void> {
+    if (this._succeeded) {
+      return;
+    }
+    this._inFlight ??= this._runOnce(work);
+
+    const inFlight = this._inFlight;
+    try {
+      await inFlight;
+    } finally {
+      // Clear only the run this caller waited on. A caller that joined the same
+      // failed run may resume after a later caller has already started a fresh
+      // one, which must not be discarded.
+      if (this._inFlight === inFlight) {
+        this._inFlight = null;
+      }
+    }
+  }
+
+  private async _runOnce(work: Work): Promise<void> {
+    await work();
+    this._succeeded = true;
+  }
+}

tests/ha/registry/device/index.test.ts

@@ -45,6 +45,23 @@ describe('DeviceRegistryManager', () => {
       expect(hass.callWS).toHaveBeenCalledTimes(1);
     });
 
+    it('should fetch once for callers that arrive while a fetch is running', async () => {
+      const testDevice = createRegistryDevice({ id: 'test' });
+
+      const hass = createHASS();
+      const manager = new DeviceRegistryManager(new DeviceCache());
+      vi.mocked(hass.callWS).mockResolvedValueOnce([testDevice]);
+
+      expect(
+        await Promise.all([
+          manager.getDevice(hass, 'test'),
+          manager.getDevice(hass, 'test'),
+        ]),
+      ).toEqual([testDevice, testDevice]);
+
+      expect(hass.callWS).toHaveBeenCalledTimes(1);
+    });
+
     it('should return null when fetch fails', async () => {
       const hass = createHASS();
       vi.mocked(hass.callWS).mockRejectedValueOnce(new Error('Fetch error'));
@@ -57,6 +74,22 @@ describe('DeviceRegistryManager', () => {
         expect.anything(),
       );
     });
+
+    it('should fetch again after a failure', async () => {
+      const testDevice = createRegistryDevice({ id: 'test' });
+
+      const hass = createHASS();
+      vi.mocked(hass.callWS)
+        .mockRejectedValueOnce(new Error('Fetch error'))
+        .mockResolvedValueOnce([testDevice]);
+
+      const manager = new DeviceRegistryManager(new DeviceCache());
+
+      expect(await manager.getDevice(hass, 'test')).toBeNull();
+      expect(await manager.getDevice(hass, 'test')).toEqual(testDevice);
+
+      expect(hass.callWS).toHaveBeenCalledTimes(2);
+    });
   });
 
   it('getMatchingDevices', async () => {

tests/ha/registry/entity/index.test.ts

@@ -101,6 +101,41 @@ describe('EntityRegistryManager', () => {
       expect(hass.callWS).toHaveBeenCalledTimes(1);
     });
 
+    it('should fetch once for callers that arrive while a fetch is running', async () => {
+      const hass = createHASS();
+      const entity = createRegistryEntity({ entity_id: 'cached' });
+      vi.mocked(hass.callWS).mockResolvedValueOnce([entity]);
+
+      const manager = new EntityRegistryManagerLive(new EntityCache());
+
+      await Promise.all([manager.fetchEntityList(hass), manager.fetchEntityList(hass)]);
+
+      expect(hass.callWS).toHaveBeenCalledTimes(1);
+      expect(await manager.getEntity(hass, 'cached')).toEqual(entity);
+    });
+
+    it('should use the first caller hass for callers that join a running fetch', async () => {
+      const entity = createRegistryEntity({ entity_id: 'cached' });
+
+      const firstHASS = createHASS();
+      vi.mocked(firstHASS.callWS).mockResolvedValueOnce([entity]);
+
+      // A later `hass` arrives mid-fetch (Home Assistant replaces the object on
+      // every state update). The joining caller uses the running fetch, so its
+      // own `hass` is never called.
+      const laterHASS = createHASS();
+
+      const manager = new EntityRegistryManagerLive(new EntityCache());
+
+      await Promise.all([
+        manager.fetchEntityList(firstHASS),
+        manager.fetchEntityList(laterHASS),
+      ]);
+
+      expect(firstHASS.callWS).toHaveBeenCalledTimes(1);
+      expect(laterHASS.callWS).not.toHaveBeenCalled();
+    });
+
     it('should log to console on error', async () => {
       const hass = createHASS();
       vi.mocked(hass.callWS).mockRejectedValueOnce(new Error('Fetch error'));
@@ -114,6 +149,22 @@ describe('EntityRegistryManager', () => {
         expect.anything(),
       );
     });
+
+    it('should fetch again after a failure', async () => {
+      const hass = createHASS();
+      const entity = createRegistryEntity({ entity_id: 'cached' });
+      vi.mocked(hass.callWS)
+        .mockRejectedValueOnce(new Error('Fetch error'))
+        .mockResolvedValueOnce([entity]);
+
+      const manager = new EntityRegistryManagerLive(new EntityCache());
+
+      await manager.fetchEntityList(hass);
+      await manager.fetchEntityList(hass);
+
+      expect(hass.callWS).toHaveBeenCalledTimes(2);
+      expect(await manager.getEntity(hass, 'cached')).toEqual(entity);
+    });
   });
 
   it('getMatchingEntities', async () => {

tests/utils/concurrency/once-runner.test.ts

@@ -0,0 +1,110 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { OnceRunner, type Work } from '../../../src/utils/concurrency/once-runner';
+
+describe('OnceRunner', () => {
+  const createDeferredWork = (): {
+    work: Work;
+    resolveAll: () => void;
+    rejectAll: (error: Error) => void;
+    callCount: () => number;
+  } => {
+    const resolvers: (() => void)[] = [];
+    const rejecters: ((error: Error) => void)[] = [];
+    const work = vi.fn().mockImplementation(
+      () =>
+        new Promise<void>((resolve, reject) => {
+          resolvers.push(resolve);
+          rejecters.push(reject);
+        }),
+    );
+    return {
+      work,
+      resolveAll: () => resolvers.forEach((resolve) => resolve()),
+      rejectAll: (error: Error) => rejecters.forEach((reject) => reject(error)),
+      callCount: () => work.mock.calls.length,
+    };
+  };
+
+  it('should share a single run with callers that arrive while it is running', async () => {
+    const runner = new OnceRunner();
+    const { work, resolveAll, callCount } = createDeferredWork();
+
+    const first = runner.run(work);
+    const second = runner.run(work);
+
+    expect(callCount()).toBe(1);
+
+    resolveAll();
+    await Promise.all([first, second]);
+
+    expect(callCount()).toBe(1);
+  });
+
+  it('should discard the work of a caller that joins a running run', async () => {
+    const runner = new OnceRunner();
+    const { work: runningWork, resolveAll, callCount } = createDeferredWork();
+    const joiningWork = vi.fn().mockResolvedValue(undefined);
+
+    const first = runner.run(runningWork);
+    const second = runner.run(joiningWork);
+
+    resolveAll();
+    await Promise.all([first, second]);
+
+    expect(callCount()).toBe(1);
+    expect(joiningWork).not.toHaveBeenCalled();
+  });
+
+  it('should not run the work again after it has succeeded', async () => {
+    const runner = new OnceRunner();
+    const work = vi.fn().mockResolvedValue(undefined);
+
+    await runner.run(work);
+    await runner.run(work);
+
+    expect(work).toHaveBeenCalledTimes(1);
+  });
+
+  it('should reject every caller waiting on a failed run', async () => {
+    const runner = new OnceRunner();
+    const { work, rejectAll } = createDeferredWork();
+    const error = new Error('failed');
+
+    const first = runner.run(work);
+    const second = runner.run(work);
+
+    rejectAll(error);
+
+    await expect(first).rejects.toThrow(error);
+    await expect(second).rejects.toThrow(error);
+  });
+
+  it('should run the work again after it throws synchronously', async () => {
+    const runner = new OnceRunner();
+    const work = vi
+      .fn()
+      .mockImplementationOnce(() => {
+        throw new Error('failed');
+      })
+      .mockResolvedValueOnce(undefined);
+
+    await expect(runner.run(work)).rejects.toThrow('failed');
+    await runner.run(work);
+
+    expect(work).toHaveBeenCalledTimes(2);
+  });
+
+  it('should run the work again after a failure', async () => {
+    const runner = new OnceRunner();
+    const work = vi
+      .fn()
+      .mockRejectedValueOnce(new Error('failed'))
+      .mockResolvedValueOnce(undefined);
+
+    await expect(runner.run(work)).rejects.toThrow('failed');
+    await runner.run(work);
+
+    expect(work).toHaveBeenCalledTimes(2);
+  });
+});