Grid LRU caching with Looker size estimates

app/packages/core/src/components/Grid/Grid.tsx

@@ -71,8 +71,9 @@ function Grid() {
       rowAspectRatioThreshold: threshold,
       get: (next) => page(next),
       render: (id, element, dimensions, zooming) => {
-        if (lookerStore.has(id.description)) {
-          lookerStore.get(id.description)?.attach(element, dimensions);
+        const cached = lookerStore.get(id.description);
+        if (cached) {
+          cached?.attach(element, dimensions);
           return;
         }
 

app/packages/core/src/components/Grid/useLookerCache.test.ts

@@ -0,0 +1,47 @@
+import { act, renderHook } from "@testing-library/react-hooks";
+import { describe, expect, it } from "vitest";
+import useLookerCache from "./useLookerCache";
+
+class Looker extends EventTarget {
+  destroy = () => undefined;
+  getSizeBytesEstimate = () => 1;
+}
+
+describe("useLookerCache", () => {
+  it("assert intermediate loading", () => {
+    const { result, rerender } = renderHook(
+      (reset) => useLookerCache<Looker>(reset),
+      {
+        initialProps: "one",
+      }
+    );
+    expect(result.current.loadedSize()).toBe(0);
+    expect(result.current.loadingSize()).toBe(0);
+
+    act(() => {
+      const looker = new Looker();
+      result.current.set("one", looker);
+    });
+
+    expect(result.current.loadingSize()).toBe(1);
+    expect(result.current.loadedSize()).toBe(0);
+    expect(result.current.sizeEstimate()).toBe(0);
+
+    act(() => {
+      const looker = result.current.get("one");
+      if (!looker) {
+        throw new Error("looker is missing");
+      }
+
+      looker.dispatchEvent(new Event("load"));
+    });
+    expect(result.current.loadingSize()).toBe(0);
+    expect(result.current.loadedSize()).toBe(1);
+    expect(result.current.sizeEstimate()).toBe(1);
+
+    rerender("two");
+    expect(result.current.loadedSize()).toBe(0);
+    expect(result.current.loadingSize()).toBe(0);
+    expect(result.current.sizeEstimate()).toBe(0);
+  });
+});

app/packages/core/src/components/Grid/useLookerCache.ts

@@ -0,0 +1,63 @@
+import * as fos from "@fiftyone/state";
+import { LRUCache } from "lru-cache";
+import { useMemo } from "react";
+
+const MAX_LRU_CACHE_ITEMS = 510;
+const MAX_LRU_CACHE_SIZE = 1e9;
+
+interface Lookers extends EventTarget {
+  destroy: () => void;
+  getSizeBytesEstimate: () => number;
+}
+
+export default function useLookerCache<
+  T extends Lookers | fos.Lookers = fos.Lookers
+>(reset: string) {
+  return useMemo(() => {
+    /** CLEAR CACHE WHEN reset CHANGES */
+    reset;
+    /** CLEAR CACHE WHEN reset CHANGES */
+
+    const loaded = new LRUCache<string, T>({
+      dispose: (looker) => looker.destroy(),
+      max: MAX_LRU_CACHE_ITEMS,
+      maxSize: MAX_LRU_CACHE_SIZE,
+      noDisposeOnSet: true,
+      sizeCalculation: (looker) => looker.getSizeBytesEstimate(),
+      updateAgeOnGet: true,
+    });
+
+    // an intermediate mapping until the "load" event
+    // "load" must occur before requesting the size bytes estimate
+    const loading = new Map<string, T>();
+
+    return {
+      delete: (key: string) => {
+        loading.delete(key);
+        loaded.delete(key);
+      },
+      get: (key: string) => loaded.get(key) ?? loading.get(key),
+      keys: function* () {
+        for (const it of loading.keys()) {
+          yield it;
+        }
+        for (const it of loaded.keys()) {
+          yield it;
+        }
+      },
+      loadingSize: () => loading.size,
+      loadedSize: () => loaded.size,
+      set: (key: string, looker: T) => {
+        const onReady = () => {
+          loaded.set(key, looker);
+          loading.delete(key);
+          looker.removeEventListener("load", onReady);
+        };
+
+        looker.addEventListener("load", onReady);
+        loading.set(key, looker);
+      },
+      sizeEstimate: () => loaded.calculatedSize,
+    };
+  }, [reset]);
+}

app/packages/core/src/components/Grid/useRefreshers.ts

@@ -1,12 +1,10 @@
 import { subscribe } from "@fiftyone/relay";
 import * as fos from "@fiftyone/state";
-import { LRUCache } from "lru-cache";
 import { useEffect, useMemo } from "react";
 import uuid from "react-uuid";
 import { useRecoilValue } from "recoil";
 import { gridAt, gridOffset, gridPage } from "./recoil";
-
-const MAX_LRU_CACHE_ITEMS = 510;
+import useLookerCache from "./useLookerCache";
 
 export default function useRefreshers() {
   const cropToContent = useRecoilValue(fos.cropToContent(false));
@@ -77,22 +75,8 @@ export default function useRefreshers() {
     []
   );
 
-  const lookerStore = useMemo(() => {
-    /** LOOKER STORE REFRESHER */
-    reset;
-    /** LOOKER STORE REFRESHER */
-
-    return new LRUCache<string, fos.Lookers>({
-      dispose: (looker) => {
-        looker.destroy();
-      },
-      max: MAX_LRU_CACHE_ITEMS,
-      noDisposeOnSet: true,
-    });
-  }, [reset]);
-
   return {
-    lookerStore,
+    lookerStore: useLookerCache(reset),
     pageReset,
     reset,
   };

app/packages/core/src/components/Grid/useSelect.ts

@@ -1,13 +1,13 @@
 import type Spotlight from "@fiftyone/spotlight";
 import * as fos from "@fiftyone/state";
-import type { LRUCache } from "lru-cache";
 import { useEffect } from "react";
 import { useRecoilValue } from "recoil";
+import type useRefreshers from "./useRefreshers";
 
 export default function useSelect(
   getFontSize: () => number,
   options: ReturnType<typeof fos.useLookerOptions>,
-  store: LRUCache<string, fos.Lookers>,
+  store: ReturnType<typeof useRefreshers>["lookerStore"],
   spotlight?: Spotlight<number, fos.Sample>
 ) {
   const { init, deferred } = fos.useDeferrer();

app/packages/looker/src/lookers/abstract.ts

@@ -191,7 +191,24 @@ export abstract class AbstractLooker<
     return this.sampleOverlays;
   }
 
-  protected dispatchEvent(eventType: string, detail: any): void {
+  getSizeBytesEstimate() {
+    let size = 1;
+    if (this.state.dimensions && !this.state.error) {
+      const [w, h] = this.state.dimensions;
+      size += w * h * 4;
+    }
+
+    if (!this.sampleOverlays?.length) {
+      return size;
+    }
+
+    for (let index = 0; index < this.sampleOverlays.length; index++) {
+      size += this.sampleOverlays[index].getSizeBytes();
+    }
+    return size;
+  }
+
+  dispatchEvent(eventType: string, detail: any): void {
     if (detail instanceof ErrorEvent) {
       this.updater({ error: detail.error });
       return;
@@ -208,11 +225,18 @@ export abstract class AbstractLooker<
   }
 
   protected dispatchImpliedEvents(
-    { options: prevOtions }: Readonly<State>,
-    { options }: Readonly<State>
+    previous: Readonly<State>,
+    next: Readonly<State>
   ): void {
-    if (options.showJSON !== prevOtions.showJSON) {
-      this.dispatchEvent("options", { showJSON: options.showJSON });
+    if (previous.options.showJSON !== next.options.showJSON) {
+      this.dispatchEvent("options", { showJSON: next.options.showJSON });
+    }
+
+    const wasLoaded = previous.overlaysPrepared && previous.dimensions;
+    const isLoaded = next.overlaysPrepared && next.dimensions;
+
+    if (!wasLoaded && isLoaded) {
+      this.dispatchEvent("load", undefined);
     }
   }
 

app/packages/looker/src/lookers/video.ts

@@ -158,6 +158,19 @@ export class VideoLooker extends AbstractLooker<VideoState, VideoSample> {
     };
   }
 
+  getSizeBytesEstimate() {
+    let size = super.getSizeBytesEstimate();
+    if (!this.firstFrame.overlays?.length) {
+      return size;
+    }
+
+    for (let index = 0; index < this.firstFrame.overlays.length; index++) {
+      size += this.firstFrame.overlays[index].getSizeBytes();
+    }
+
+    return size;
+  }
+
   hasDefaultZoom(state: VideoState, overlays: Overlay<VideoState>[]): boolean {
     const pan = [0, 0];
     const scale = 1;