perf(observation): reduce document processing scheduling overhead

apps/extension/src/tools/vom/__tests__/facts.test.ts

@@ -1,5 +1,6 @@
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { OVERLAY_HOST_MARKER_ATTR } from "@/lib/overlay-bridge";
+import { createCaptureCheckpoint } from "../capture-abort";
 import { buildDocumentIndex, type DecodedNode } from "../facts";
 import { decodeDocument, REQUESTED_STYLES } from "../snapshot";
 
@@ -90,7 +91,23 @@ describe("document facts", () => {
     expect(nodes[0]).not.toHaveProperty("excluded");
   });
 
-  it("yields during large inputs so cancellation interrupts indexing", async () => {
+  it("keeps short blocks synchronous and checks cancellation after yielding", async () => {
+    const clock = vi.spyOn(performance, "now").mockReturnValue(0);
+    try {
+      const controller = new AbortController();
+      const checkpoint = createCaptureCheckpoint(controller.signal);
+      for (let i = 0; i < 20; i++) expect(checkpoint()).toBeUndefined();
+      clock.mockReturnValue(9);
+      const pending = checkpoint();
+      expect(pending).toBeInstanceOf(Promise);
+      controller.abort();
+      await expect(pending).rejects.toMatchObject({ name: "AbortError" });
+    } finally {
+      clock.mockRestore();
+    }
+  });
+
+  it("checks cancellation at bounded indexing blocks", async () => {
     const controller = new AbortController();
     const input = Array.from({ length: 100_000 }, (_, i) => node(i, i ? i - 1 : null));
     let reads = 0;

apps/extension/src/tools/vom/capture-abort.ts

@@ -16,17 +16,25 @@ export function throwIfAborted(signal: AbortSignal | undefined): void {
   if (signal?.aborted) throw captureAbortError();
 }
 
-/** Called only at bounded block boundaries, not once per node. */
+/** Check at bounded work blocks; yield only after a processing slice is spent.
-export async function captureCheckpoint(signal?: AbortSignal): Promise<void> {
+ * Each caller owns its deadline, so unrelated captures do not share state. */
-  throwIfAborted(signal);
+export function createCaptureCheckpoint(signal?: AbortSignal): () => Promise<void> | undefined {
-  await new Promise<void>((resolve) => {
-    const channel = new MessageChannel();
-    channel.port1.onmessage = () => {
-      channel.port1.close();
-      channel.port2.close();
-      resolve();
-    };
-    channel.port2.postMessage(null);
-  });
   throwIfAborted(signal);
+  let deadline = performance.now() + 8;
+  return () => {
+    throwIfAborted(signal);
+    if (performance.now() < deadline) return;
+    return new Promise<void>((resolve) => {
+      const channel = new MessageChannel();
+      channel.port1.onmessage = () => {
+        channel.port1.close();
+        channel.port2.close();
+        resolve();
+      };
+      channel.port2.postMessage(null);
+    }).then(() => {
+      throwIfAborted(signal);
+      deadline = performance.now() + 8;
+    });
+  };
 }

apps/extension/src/tools/vom/facts.ts

@@ -1,5 +1,5 @@
 import { isOverlayHostNode } from "@/lib/overlay-bridge";
-import { captureCheckpoint } from "./capture-abort";
+import { createCaptureCheckpoint } from "./capture-abort";
 import type { CapturedNode } from "./capture-types";
 
 /** Bounds retain their snapshot document CSS units until normalization. */
@@ -34,24 +34,34 @@ export async function buildDocumentIndex<T extends DecodedNode>(
   input: readonly T[],
   signal?: AbortSignal,
 ): Promise<DocumentIndex<T>> {
+  const checkpoint = createCaptureCheckpoint(signal);
   const nodes = new Map<number, T>();
   const overlayByNode = new Map<number, boolean>();
   const excludedBackendNodeIds = new Set<number>();
   for (let i = 0; i < input.length; i++) {
-    if (i % 256 === 0) await captureCheckpoint(signal);
+    if (i % 256 === 0) {
+      const pending = checkpoint();
+      if (pending) await pending;
+    }
     const node = input[i];
     nodes.set(node.backendNodeId, node);
   }
   let work = 0;
   for (const node of input) {
-    if (work++ % 256 === 0) await captureCheckpoint(signal);
+    if (work++ % 256 === 0) {
+      const pending = checkpoint();
+      if (pending) await pending;
+    }
     if (overlayByNode.has(node.backendNodeId)) continue;
     const path: DecodedNode[] = [];
     const visiting = new Set<number>();
     let current: DecodedNode | undefined = node;
     let overlay = false;
     while (current && !overlayByNode.has(current.backendNodeId)) {
-      if (work++ % 256 === 0) await captureCheckpoint(signal);
+      if (work++ % 256 === 0) {
+        const pending = checkpoint();
+        if (pending) await pending;
+      }
       if (visiting.has(current.backendNodeId)) {
         // Only an overlay inside the cycle may mark the cycle and its descendants.
         overlay = path
@@ -70,7 +80,10 @@ export async function buildDocumentIndex<T extends DecodedNode>(
     if (current && overlayByNode.has(current.backendNodeId))
       overlay = overlayByNode.get(current.backendNodeId)!;
     for (let i = path.length - 1; i >= 0; i--) {
-      if (work++ % 256 === 0) await captureCheckpoint(signal);
+      if (work++ % 256 === 0) {
+        const pending = checkpoint();
+        if (pending) await pending;
+      }
       const item = path[i];
       overlay = overlay || isOverlayHostNode(item.tag, Object.keys(item.attrs));
       overlayByNode.set(item.backendNodeId, overlay);

apps/extension/src/tools/vom/frame-document.ts

@@ -4,7 +4,7 @@ import {
   type CdpTarget,
   cdpTargetKey,
 } from "@/browser-driver/frame-graph";
-import { captureCheckpoint } from "./capture-abort";
+import { createCaptureCheckpoint } from "./capture-abort";
 import type { CapturedNode } from "./capture-types";
 
 export interface FrameOwnedAxNode {
@@ -88,6 +88,7 @@ export async function buildFrameDocuments<T extends FrameOwnedAxNode>(
   captured: FrameDomInput,
   signal?: AbortSignal,
 ): Promise<FrameDocument<T>[]> {
+  const checkpoint = createCaptureCheckpoint(signal);
   let work = 0;
   const frames = frameList(graph, batches, captured);
   const frameById = new Map(frames.map((frame) => [frame.frameId, frame]));
@@ -97,7 +98,10 @@ export async function buildFrameDocuments<T extends FrameOwnedAxNode>(
   const backendOwner = new Map<string, string>();
   for (const frame of frames) {
     for (const node of domNodesForFrame(frame.frameId)) {
-      if (work++ % 256 === 0) await captureCheckpoint(signal);
+      if (work++ % 256 === 0) {
+        const pending = checkpoint();
+        if (pending) await pending;
+      }
       backendOwner.set(targetBackendKey(frame.target, node.backendNodeId), frame.frameId);
     }
   }
@@ -106,17 +110,23 @@ export async function buildFrameDocuments<T extends FrameOwnedAxNode>(
   for (const batch of batches) {
     const nodeById = new Map<string, T>();
     for (const node of batch.nodes) {
-      if (work++ % 256 === 0) await captureCheckpoint(signal);
+      if (work++ % 256 === 0) {
+        const pending = checkpoint();
+        if (pending) await pending;
+      }
       nodeById.set(node.nodeId, node);
     }
     const ownershipByNodeId = new Map<string, Ownership>();
-    const resolveOwnership = async (node: T): Promise<Ownership> => {
+    for (const node of batch.nodes) {
       const path: T[] = [];
       const visiting = new Set<string>();
       let current: T | undefined = node;
       let ownership: Ownership = { frameId: batch.frame.frameId, strength: 1 };
       while (current) {
-        if (work++ % 256 === 0) await captureCheckpoint(signal);
+        if (work++ % 256 === 0) {
+          const pending = checkpoint();
+          if (pending) await pending;
+        }
         const cached = ownershipByNodeId.get(current.nodeId);
         if (cached) {
           ownership = cached;
@@ -142,17 +152,14 @@ export async function buildFrameDocuments<T extends FrameOwnedAxNode>(
         current = current.parentId ? nodeById.get(current.parentId) : undefined;
       }
       for (let i = path.length - 1; i >= 0; i--) {
-        if (work++ % 256 === 0) await captureCheckpoint(signal);
+        if (work++ % 256 === 0) {
+          const pending = checkpoint();
+          if (pending) await pending;
+        }
         ownership = { frameId: ownership.frameId, strength: Math.min(2, ownership.strength) };
         ownershipByNodeId.set(path[i].nodeId, ownership);
       }
       if (!ownershipByNodeId.has(node.nodeId)) ownershipByNodeId.set(node.nodeId, ownership);
-      return ownership;
-    };
-
-    for (const node of batch.nodes) {
-      if (work++ % 256 === 0) await captureCheckpoint(signal);
-      const ownership = await resolveOwnership(node);
       const key = targetNodeKey(batch.frame.target, node.nodeId);
       const existing = candidates.get(key);
       if (!existing || ownership.strength > existing.ownership.strength) {
@@ -166,7 +173,10 @@ export async function buildFrameDocuments<T extends FrameOwnedAxNode>(
   );
   const axNodesByFrame = new Map<string, T[]>();
   for (const candidate of candidates.values()) {
-    if (work++ % 256 === 0) await captureCheckpoint(signal);
+    if (work++ % 256 === 0) {
+      const pending = checkpoint();
+      if (pending) await pending;
+    }
     const { node, target, ownership } = candidate;
     const parentFrameId = node.parentId
       ? ownershipByTargetNode.get(targetNodeKey(target, node.parentId))

apps/extension/src/tools/vom/normalize.ts

@@ -4,7 +4,7 @@ import {
   projectSnapshotRect,
   snapshotViewportRect,
 } from "../geometry/coordinate-types";
-import { captureCheckpoint } from "./capture-abort";
+import { createCaptureCheckpoint } from "./capture-abort";
 import { buildDocumentIndex, type DocumentIndex, type NodeFacts } from "./facts";
 import { decodeDocument, type SnapshotDocument } from "./snapshot";
 
@@ -30,9 +30,13 @@ export async function normalizeDocument(
   signal?: AbortSignal,
 ): Promise<NormalizedDocument> {
   const decoded = await decodeDocument(doc, strings, signal);
+  const checkpoint = createCaptureCheckpoint(signal);
   const nodes: NodeFacts[] = [];
   for (let i = 0; i < decoded.nodes.length; i++) {
-    if (i % 256 === 0) await captureCheckpoint(signal);
+    if (i % 256 === 0) {
+      const pending = checkpoint();
+      if (pending) await pending;
+    }
     const node = decoded.nodes[i];
     const input = snapshotViewportRect(
       node.layout?.bounds ?? [],

apps/extension/src/tools/vom/snapshot.ts

@@ -1,4 +1,4 @@
-import { captureCheckpoint } from "./capture-abort";
+import { createCaptureCheckpoint } from "./capture-abort";
 import type { CapturedNode } from "./capture-types";
 import type { DecodedDocument, DecodedNode } from "./facts";
 export const REQUESTED_STYLES = [
@@ -102,6 +102,7 @@ export async function decodeDocument(
   strings: string[],
   signal?: AbortSignal,
 ): Promise<DecodedDocument> {
+  const checkpoint = createCaptureCheckpoint(signal);
   const dn = doc.nodes;
   const dl = doc.layout;
   if (!dn?.backendNodeId) {
@@ -111,7 +112,10 @@ export async function decodeDocument(
   const count = dn.backendNodeId.length;
   const layoutByNode = new Map<number, number>();
   for (let i = 0; i < (dl?.nodeIndex?.length ?? 0); i++) {
-    if (i % 256 === 0) await captureCheckpoint(signal);
+    if (i % 256 === 0) {
+      const pending = checkpoint();
+      if (pending) await pending;
+    }
     layoutByNode.set(dl!.nodeIndex![i], i);
   }
 
@@ -126,7 +130,10 @@ export async function decodeDocument(
   const nodeTextContent = new Map<number, string[]>();
   if (dn.nodeValue) {
     for (let n = 0; n < count; n++) {
-      if (n % 256 === 0) await captureCheckpoint(signal);
+      if (n % 256 === 0) {
+        const pending = checkpoint();
+        if (pending) await pending;
+      }
       const nvIdx = dn.nodeValue[n] ?? -1;
       if (nvIdx < 0) continue;
       const text = str(strings, nvIdx).trim();
@@ -143,7 +150,10 @@ export async function decodeDocument(
   // Decode fields without assigning a coordinate projection or semantic policy.
   const nodes: DecodedNode[] = [];
   for (let n = 0; n < count; n++) {
-    if (n % 256 === 0) await captureCheckpoint(signal);
+    if (n % 256 === 0) {
+      const pending = checkpoint();
+      if (pending) await pending;
+    }
     const backendNodeId = dn.backendNodeId[n];
     const parentIdx = dn.parentIndex?.[n] ?? -1;
     const parentBackendNodeId = parentIdx >= 0 ? (dn.backendNodeId[parentIdx] ?? null) : null;