improve overlay rendering performance

app/packages/looker/src/elements/common/error.ts

@@ -103,12 +103,11 @@ export class ErrorElement<State extends BaseState> extends BaseElement<State> {
       }
     }
 
+    if (!error && this.errorElement) {
+      this.errorElement.remove();
+      this.errorElement = null;
+    }
+
     return this.errorElement;
   }
 }
-
-const onClick = (href) => {
-  let openExternal;
-
-  return null;
-};

app/packages/looker/src/elements/image.ts

@@ -6,22 +6,48 @@ import type { ImageState } from "../state";
 import type { Events } from "./base";
 import { BaseElement } from "./base";
 
+const MAX_IMAGE_LOAD_RETRIES = 10;
+
 export class ImageElement extends BaseElement<ImageState, HTMLImageElement> {
   private src = "";
-  private imageSource: HTMLImageElement;
+  protected imageSource: HTMLImageElement;
+
+  private retryCount = 0;
+  private timeoutId: number | null = null;
 
   getEvents(): Events<ImageState> {
     return {
       load: ({ update }) => {
+        if (this.timeoutId !== null) {
+          window.clearTimeout(this.timeoutId);
+          this.timeoutId = null;
+        }
+        this.retryCount = 0;
+
         this.imageSource = this.element;
 
         update({
           loaded: true,
+          error: false,
           dimensions: [this.element.naturalWidth, this.element.naturalHeight],
         });
       },
       error: ({ update }) => {
         update({ error: true, dimensions: [512, 512], loaded: true });
+        // sometimes image loading fails because of insufficient resources
+        // we'll want to try again in those cases
+        if (this.retryCount < MAX_IMAGE_LOAD_RETRIES) {
+          // schedule a retry after a delay
+          if (this.timeoutId !== null) {
+            window.clearTimeout(this.timeoutId);
+          }
+          this.timeoutId = window.setTimeout(() => {
+            this.retryCount += 1;
+            const retrySrc = `${this.src}`;
+            this.element.setAttribute("src", retrySrc);
+            // linear backoff
+          }, 1000 * this.retryCount);
+        }
       },
     };
   }
@@ -36,10 +62,13 @@ export class ImageElement extends BaseElement<ImageState, HTMLImageElement> {
   renderSelf({ config: { src } }: Readonly<ImageState>) {
     if (this.src !== src) {
       this.src = src;
-
+      this.retryCount = 0;
+      if (this.timeoutId !== null) {
+        window.clearTimeout(this.timeoutId);
+        this.timeoutId = null;
+      }
       this.element.setAttribute("src", src);
     }
-
     return null;
   }
 }

app/packages/looker/src/worker/decorated-fetch.test.ts

@@ -0,0 +1,80 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { fetchWithLinearBackoff } from "./decorated-fetch";
+
+describe("fetchWithLinearBackoff", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.useRealTimers();
+  });
+
+  it("should return response when fetch succeeds on first try", async () => {
+    const mockResponse = new Response("Success", { status: 200 });
+    global.fetch = vi.fn().mockResolvedValue(mockResponse);
+
+    const response = await fetchWithLinearBackoff("http://fiftyone.ai");
+
+    expect(response).toBe(mockResponse);
+    expect(global.fetch).toHaveBeenCalledTimes(1);
+    expect(global.fetch).toHaveBeenCalledWith("http://fiftyone.ai");
+  });
+
+  it("should retry when fetch fails and eventually succeed", async () => {
+    const mockResponse = new Response("Success", { status: 200 });
+    global.fetch = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("Network Error"))
+      .mockResolvedValue(mockResponse);
+
+    const response = await fetchWithLinearBackoff("http://fiftyone.ai");
+
+    expect(response).toBe(mockResponse);
+    expect(global.fetch).toHaveBeenCalledTimes(2);
+  });
+
+  it("should throw an error after max retries when fetch fails every time", async () => {
+    global.fetch = vi.fn().mockRejectedValue(new Error("Network Error"));
+
+    await expect(
+      fetchWithLinearBackoff("http://fiftyone.ai", 3, 10)
+    ).rejects.toThrowError(new RegExp("Max retries for fetch reached"));
+
+    expect(global.fetch).toHaveBeenCalledTimes(3);
+  });
+
+  it("should throw an error when response is not ok", async () => {
+    const mockResponse = new Response("Not Found", { status: 404 });
+    global.fetch = vi.fn().mockResolvedValue(mockResponse);
+
+    await expect(
+      fetchWithLinearBackoff("http://fiftyone.ai", 5, 10)
+    ).rejects.toThrow("HTTP error: 404");
+
+    expect(global.fetch).toHaveBeenCalledTimes(5);
+  });
+
+  it("should apply linear backoff between retries", async () => {
+    const mockResponse = new Response("Success", { status: 200 });
+    global.fetch = vi
+      .fn()
+      .mockRejectedValueOnce(new Error("Network Error"))
+      .mockRejectedValueOnce(new Error("Network Error"))
+      .mockResolvedValue(mockResponse);
+
+    vi.useFakeTimers();
+
+    const fetchPromise = fetchWithLinearBackoff("http://fiftyone.ai", 5, 10);
+
+    // advance timers to simulate delays
+    // after first delay
+    await vi.advanceTimersByTimeAsync(100);
+    // after scond delay
+    await vi.advanceTimersByTimeAsync(200);
+
+    const response = await fetchPromise;
+
+    expect(response).toBe(mockResponse);
+    expect(global.fetch).toHaveBeenCalledTimes(3);
+
+    vi.useRealTimers();
+  });
+});

app/packages/looker/src/worker/decorated-fetch.ts

@@ -0,0 +1,25 @@
+const DEFAULT_MAX_RETRIES = 10;
+const DEFAULT_BASE_DELAY = 200;
+
+export const fetchWithLinearBackoff = async (
+  url: string,
+  retries = DEFAULT_MAX_RETRIES,
+  delay = DEFAULT_BASE_DELAY
+) => {
+  for (let i = 0; i < retries; i++) {
+    try {
+      const response = await fetch(url);
+      if (!response.ok) throw new Error(`HTTP error: ${response.status}`);
+      return response;
+    } catch (e) {
+      if (i < retries - 1) {
+        await new Promise((resolve) => setTimeout(resolve, delay * (i + 1)));
+      } else {
+        throw new Error(
+          "Max retries for fetch reached (linear backoff), error: " + e
+        );
+      }
+    }
+  }
+  return null;
+};

app/packages/looker/src/worker/index.ts

@@ -30,6 +30,7 @@ import {
   Sample,
 } from "../state";
 import { decodeWithCanvas } from "./canvas-decoder";
+import { fetchWithLinearBackoff } from "./decorated-fetch";
 import { DeserializerFactory } from "./deserializer";
 import { PainterFactory } from "./painter";
 import { mapId } from "./shared";
@@ -107,11 +108,12 @@ const imputeOverlayFromPath = async (
   colorscale: Colorscale,
   buffers: ArrayBuffer[],
   sources: { [path: string]: string },
-  cls: string
+  cls: string,
+  maskPathDecodingPromises: Promise<void>[] = []
 ) => {
   // handle all list types here
   if (cls === DETECTIONS) {
-    const promises = [];
+    const promises: Promise<void>[] = [];
     for (const detection of label.detections) {
       promises.push(
         imputeOverlayFromPath(
@@ -126,10 +128,7 @@ const imputeOverlayFromPath = async (
         )
       );
     }
-    // if some paths fail to load, it's okay, we can still proceed
-    // hence we use `allSettled` instead of `all`
-    await Promise.allSettled(promises);
-    return;
+    maskPathDecodingPromises.push(...promises);
   }
 
   // overlay path is in `map_path` property for heatmap, or else, it's in `mask_path` property (for segmentation or detection)
@@ -157,14 +156,17 @@ const imputeOverlayFromPath = async (
     baseUrl = overlayImageUrl.split("?")[0];
   }
 
-  const overlayImageBuffer: Blob = await getFetchFunction()(
-    "GET",
-    overlayImageUrl,
-    null,
-    "blob"
-  );
+  let overlayImageBlob: Blob;
+  try {
+    const overlayImageFetchResponse = await fetchWithLinearBackoff(baseUrl);
+    overlayImageBlob = await overlayImageFetchResponse.blob();
+  } catch (e) {
+    console.error(e);
+    // skip decoding if fetch fails altogether
+    return;
+  }
 
-  const overlayMask = await decodeWithCanvas(overlayImageBuffer);
+  const overlayMask = await decodeWithCanvas(overlayImageBlob);
   const [overlayHeight, overlayWidth] = overlayMask.shape;
 
   // set the `mask` property for this label
@@ -190,8 +192,11 @@ const processLabels = async (
   schema: Schema
 ): Promise<ArrayBuffer[]> => {
   const buffers: ArrayBuffer[] = [];
-  const promises = [];
+  const painterPromises = [];
+
+  const maskPathDecodingPromises = [];
 
+  // mask deserialization / mask_path decoding loop
   for (const field in sample) {
     let labels = sample[field];
     if (!Array.isArray(labels)) {
@@ -205,20 +210,19 @@ const processLabels = async (
       }
 
       if (DENSE_LABELS.has(cls)) {
-        try {
-          await imputeOverlayFromPath(
+        maskPathDecodingPromises.push(
+          imputeOverlayFromPath(
             `${prefix || ""}${field}`,
             label,
             coloring,
             customizeColorSetting,
             colorscale,
             buffers,
             sources,
-            cls
-          );
-        } catch (e) {
-          console.error("Couldn't decode overlay image from disk: ", e);
-        }
+            cls,
+            maskPathDecodingPromises
+          )
+        );
       }
 
       if (cls in DeserializerFactory) {
@@ -249,9 +253,25 @@ const processLabels = async (
           mapId(label);
         }
       }
+    }
+  }
+
+  await Promise.allSettled(maskPathDecodingPromises);
 
+  // overlay painting loop
+  for (const field in sample) {
+    let labels = sample[field];
+    if (!Array.isArray(labels)) {
+      labels = [labels];
+    }
+    const cls = getCls(`${prefix ? prefix : ""}${field}`, schema);
+
+    for (const label of labels) {
+      if (!label) {
+        continue;
+      }
       if (painterFactory[cls]) {
-        promises.push(
+        painterPromises.push(
           painterFactory[cls](
             prefix ? prefix + field : field,
             label,
@@ -266,7 +286,7 @@ const processLabels = async (
     }
   }
 
-  return Promise.all(promises).then(() => buffers);
+  return Promise.all(painterPromises).then(() => buffers);
 };
 
 /** GLOBALS */

app/packages/utilities/src/fetch.ts

@@ -18,7 +18,7 @@ export interface FetchFunction {
     body?: A,
     result?: "json" | "blob" | "text" | "arrayBuffer" | "json-stream",
     retries?: number,
-    retryCodes?: number[] | "arrayBuffer"
+    retryCodes?: number[]
   ): Promise<R>;
 }
 
@@ -110,7 +110,7 @@ export const setFetchFunction = (
     const fetchCall = retries
       ? fetchRetry(fetch, {
           retries,
-          retryDelay: 0,
+          retryDelay: 500,
           retryOn: (attempt, error, response) => {
             if (
               (error !== null || retryCodes.includes(response.status)) &&