fix(viewer): drain WebSocket messages in frame-budgeted slices, not all at once

src/communications/viewer_connection.ts

@@ -6,10 +6,30 @@ interface ViewerConnectionOptions extends ViewerWebSocketOptions {
   onError: (error: Error) => void;
 }
 
+// A caller loading a large scene (e.g. one element per WebSocket message, no batching
+// on the send side either - see compas_threejs's Outbox) can burst thousands of
+// messages at once. options.dispatch does real work per message (protobuf decode plus
+// a Three.js scene mutation) - draining the whole burst synchronously inside one
+// onmessage callback pins the main thread for the entire burst, which is what used to
+// make the tab look/become unresponsive (and, past the browser's own patience, drop
+// the WebSocket entirely) for a large model. Spending only a slice of each animation
+// frame on drained messages, and picking up where it left off next frame, keeps every
+// frame responsive - the burst still finishes in roughly the same wall-clock time, it
+// just no longer blocks anything else (rendering, input, the connection's own
+// liveness) while doing it.
+const FRAME_BUDGET_MS = 8;
+
 export class ViewerConnection {
   private socket: WebSocket | null = null;
   private retryTimer: ReturnType<typeof setTimeout> | null = null;
   private stopped = true;
+  // Queue + head index rather than Array.shift() per message - shift() is O(n), which
+  // would make draining a multi-thousand-message burst O(n^2). Reset to empty once
+  // fully drained (see drain()) so a long-lived connection doesn't hold onto an
+  // ever-growing backing array across many bursts over its lifetime.
+  private readonly inbox: Uint8Array[] = [];
+  private inboxHead = 0;
+  private drainHandle: number | null = null;
 
   constructor(private readonly options: ViewerConnectionOptions) {}
 
@@ -41,6 +61,12 @@ export class ViewerConnection {
       clearTimeout(this.retryTimer);
       this.retryTimer = null;
     }
+    if (this.drainHandle !== null) {
+      cancelAnimationFrame(this.drainHandle);
+      this.drainHandle = null;
+    }
+    this.inbox.length = 0;
+    this.inboxHead = 0;
     const socket = this.socket;
     this.socket = null;
     if (socket) {
@@ -57,7 +83,8 @@ export class ViewerConnection {
     socket.binaryType = "arraybuffer";
     socket.onmessage = (event: MessageEvent) => {
       if (event.data instanceof ArrayBuffer) {
-        this.options.dispatch(new Uint8Array(event.data));
+        this.inbox.push(new Uint8Array(event.data));
+        this.scheduleDrain();
       }
     };
     socket.onerror = () => {
@@ -71,6 +98,46 @@ export class ViewerConnection {
     };
   }
 
+  private scheduleDrain(): void {
+    if (this.drainHandle !== null) return;
+    this.drainHandle = requestAnimationFrame(() => this.drain());
+  }
+
+  /**
+   * Dispatches queued messages, in order, for up to FRAME_BUDGET_MS before yielding
+   * back to the browser - then reschedules itself for the next frame if anything is
+   * still queued. See FRAME_BUDGET_MS's own docstring for why this exists.
+   */
+  private drain(): void {
+    this.drainHandle = null;
+    const deadline = performance.now() + FRAME_BUDGET_MS;
+    while (this.inboxHead < this.inbox.length && performance.now() < deadline) {
+      const message = this.inbox[this.inboxHead]!;
+      // Advance past this message BEFORE dispatching it, and catch a throw from
+      // dispatch() itself - one malformed/unsupported message must never silently
+      // strand every message queued after it (which would otherwise happen here: an
+      // uncaught throw exits this loop without rescheduling, so the drain simply never
+      // resumes). options.dispatch already reports most failures through its own
+      // onError callback instead of throwing, but this is the difference between one
+      // skipped message and the rest of a large model - including a final "stop
+      // loading" message - never arriving at all.
+      this.inboxHead += 1;
+      try {
+        this.options.dispatch(message);
+      } catch (error) {
+        this.options.onError(
+          error instanceof Error ? error : new Error(String(error)),
+        );
+      }
+    }
+    if (this.inboxHead >= this.inbox.length) {
+      this.inbox.length = 0;
+      this.inboxHead = 0;
+    } else {
+      this.scheduleDrain();
+    }
+  }
+
   private buildUrl(): string {
     const query = new URLSearchParams(window.location.search);
     const host = this.options.host ?? query.get("ws_host") ?? "127.0.0.1";

tests/viewer_connection.test.ts

@@ -0,0 +1,124 @@
+/** @vitest-environment happy-dom */
+
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+import { ViewerConnection } from "../src/communications/viewer_connection";
+
+class MockWebSocket {
+  static readonly OPEN = 1;
+  readyState = 1;
+  binaryType = "";
+  onmessage: ((event: MessageEvent) => void) | null = null;
+  onerror: (() => void) | null = null;
+  onclose: (() => void) | null = null;
+
+  constructor(public url: string) {}
+
+  send(): void {}
+  close(): void {}
+}
+
+function bytes(value: number): Uint8Array {
+  return new Uint8Array([value]);
+}
+
+/**
+ * Covers ViewerConnection's own message-batching (see its FRAME_BUDGET_MS docstring):
+ * a burst of messages arriving in one synchronous tick - the shape a large model's
+ * load produces, one WebSocket message per element with no batching on the send side
+ * either - must drain across several animation frames rather than being dispatched
+ * (protobuf-decode-plus-scene-mutation, real work per message) all in the same
+ * onmessage callback, which is what used to pin the main thread for the whole burst.
+ */
+describe("ViewerConnection message batching", () => {
+  let sockets: MockWebSocket[];
+  let rafCallbacks: FrameRequestCallback[];
+  let nowValue: number;
+
+  beforeEach(() => {
+    sockets = [];
+    rafCallbacks = [];
+    nowValue = 0;
+    vi.stubGlobal(
+      "WebSocket",
+      class extends MockWebSocket {
+        constructor(url: string) {
+          super(url);
+          sockets.push(this);
+        }
+      },
+    );
+    vi.stubGlobal("requestAnimationFrame", (cb: FrameRequestCallback) => {
+      rafCallbacks.push(cb);
+      return rafCallbacks.length;
+    });
+    vi.stubGlobal("cancelAnimationFrame", () => {});
+    vi.spyOn(performance, "now").mockImplementation(() => nowValue);
+  });
+
+  afterEach(() => {
+    vi.unstubAllGlobals();
+    vi.restoreAllMocks();
+  });
+
+  /** Runs exactly the callbacks scheduled so far, as one animation frame would. */
+  function flushOneFrame(): void {
+    const callbacks = rafCallbacks;
+    rafCallbacks = [];
+    for (const callback of callbacks) callback(nowValue);
+  }
+
+  it("drains a burst of messages across multiple frames, in order, not all at once", () => {
+    const dispatched: Uint8Array[] = [];
+    const connection = new ViewerConnection({
+      dispatch: (message) => {
+        dispatched.push(message);
+        // Each dispatch "costs" 5ms of simulated time - with an 8ms frame budget,
+        // at most two fit per frame, forcing this test's burst across several frames.
+        nowValue += 5;
+      },
+      onError: () => {},
+    });
+    connection.start();
+    const socket = sockets[0]!;
+
+    for (let i = 0; i < 5; i++) {
+      socket.onmessage?.({ data: bytes(i).buffer } as MessageEvent);
+    }
+
+    // Nothing dispatched synchronously from onmessage, and only one frame scheduled
+    // for the whole burst - not one dispatch (or one rAF) per message.
+    expect(dispatched).toHaveLength(0);
+    expect(rafCallbacks).toHaveLength(1);
+
+    flushOneFrame();
+    expect(dispatched.length).toBeGreaterThan(0);
+    expect(dispatched.length).toBeLessThan(5);
+
+    let guard = 0;
+    while (dispatched.length < 5 && guard < 10) {
+      flushOneFrame();
+      guard += 1;
+    }
+
+    expect(dispatched).toHaveLength(5);
+    expect(dispatched.map((message) => message[0])).toEqual([0, 1, 2, 3, 4]);
+  });
+
+  it("dispose() cancels a pending drain and drops whatever was still queued", () => {
+    const dispatched: Uint8Array[] = [];
+    const connection = new ViewerConnection({
+      dispatch: (message) => dispatched.push(message),
+      onError: () => {},
+    });
+    connection.start();
+    const socket = sockets[0]!;
+    socket.onmessage?.({ data: bytes(0).buffer } as MessageEvent);
+    expect(rafCallbacks).toHaveLength(1);
+
+    connection.dispose();
+    flushOneFrame();
+
+    expect(dispatched).toHaveLength(0);
+  });
+});