From cf9533f88245475402e1ba18ffaae2789d78d100 Mon Sep 17 00:00:00 2001 From: Aiden Bai Date: Mon, 3 Aug 2026 15:14:56 -0700 Subject: [PATCH] Fix unfreeze renderer scheduling (#612) * fix unfreeze renderer scheduling * harden unfreeze edge cases * refactor renderer ownership tracking * fix stale refresh across freeze sessions --- .../react-grab/e2e/freeze-updates.spec.ts | 98 ++++++++++++++++ .../react-grab/src/utils/freeze-updates.ts | 106 +++++++++++++++--- 2 files changed, 187 insertions(+), 17 deletions(-) diff --git a/packages/react-grab/e2e/freeze-updates.spec.ts b/packages/react-grab/e2e/freeze-updates.spec.ts index 33535d2f4..54f44462a 100644 --- a/packages/react-grab/e2e/freeze-updates.spec.ts +++ b/packages/react-grab/e2e/freeze-updates.spec.ts @@ -90,6 +90,104 @@ test.describe("Freeze Updates", () => { }); test.describe("Multiple Freeze/Unfreeze Cycles", () => { + test("ignores malformed renderer registrations", async ({ reactGrab }) => { + await reactGrab.page.evaluate(() => { + const devtoolsHook = Reflect.get(window, "__REACT_DEVTOOLS_GLOBAL_HOOK__"); + if (!devtoolsHook || typeof devtoolsHook !== "object") { + throw new Error("React DevTools hook is unavailable"); + } + const injectRenderer = Reflect.get(devtoolsHook, "inject"); + if (typeof injectRenderer !== "function") { + throw new Error("React DevTools renderer injection is unavailable"); + } + const incompleteRenderer = { scheduleRefresh: () => {} }; + Reflect.apply(injectRenderer, devtoolsHook, [incompleteRenderer]); + const inaccessibleRenderer = Object.create(null); + Object.defineProperty(inaccessibleRenderer, "rendererPackageName", { + get: () => { + throw new Error("Renderer metadata is inaccessible"); + }, + }); + Reflect.apply(injectRenderer, devtoolsHook, [inaccessibleRenderer]); + + window.freezeReactGrab(); + window.unfreezeReactGrab(); + }); + }); + + test("schedules only the owning renderer and skips stale updates", async ({ reactGrab }) => { + const scheduleUpdateCounts = await reactGrab.page.evaluate(async () => { + const devtoolsHook = Reflect.get(window, "__REACT_DEVTOOLS_GLOBAL_HOOK__"); + if (!devtoolsHook || typeof devtoolsHook !== "object") { + throw new Error("React DevTools hook is unavailable"); + } + const injectRenderer = Reflect.get(devtoolsHook, "inject"); + const commitFiberRoot = Reflect.get(devtoolsHook, "onCommitFiberRoot"); + if (typeof injectRenderer !== "function" || typeof commitFiberRoot !== "function") { + throw new Error("React DevTools instrumentation is unavailable"); + } + let owningRendererScheduleUpdateCount = 0; + let foreignRendererScheduleUpdateCount = 0; + const owningRenderer = { + currentDispatcherRef: null, + rendererPackageName: "react-dom-owning", + scheduleUpdate: () => { + owningRendererScheduleUpdateCount += 1; + }, + }; + const foreignRenderer = { + currentDispatcherRef: null, + rendererPackageName: "react-dom-foreign", + scheduleUpdate: () => { + foreignRendererScheduleUpdateCount += 1; + }, + }; + const owningRendererId = Reflect.apply(injectRenderer, devtoolsHook, [owningRenderer]); + Reflect.apply(injectRenderer, devtoolsHook, [foreignRenderer]); + const syntheticFiberRoot = Object.create(null); + const syntheticRootFiber = { + child: null, + return: null, + sibling: null, + stateNode: syntheticFiberRoot, + }; + Object.assign(syntheticFiberRoot, { + containerInfo: document.body, + current: syntheticRootFiber, + }); + Reflect.apply(commitFiberRoot, devtoolsHook, [owningRendererId, syntheticFiberRoot]); + + window.freezeReactGrab(); + window.unfreezeReactGrab(); + window.freezeReactGrab(); + await new Promise((resolve) => queueMicrotask(resolve)); + const staleScheduleUpdateCount = owningRendererScheduleUpdateCount; + window.unfreezeReactGrab(); + await new Promise((resolve) => queueMicrotask(resolve)); + + const scheduleUpdateCountBeforeRapidCycle = owningRendererScheduleUpdateCount; + window.freezeReactGrab(); + window.unfreezeReactGrab(); + window.freezeReactGrab(); + window.unfreezeReactGrab(); + await new Promise((resolve) => queueMicrotask(resolve)); + return { + foreignRendererScheduleUpdateCount, + owningRendererScheduleUpdateCount, + rapidCycleScheduleUpdateCount: + owningRendererScheduleUpdateCount - scheduleUpdateCountBeforeRapidCycle, + staleScheduleUpdateCount, + }; + }); + + expect(scheduleUpdateCounts).toEqual({ + foreignRendererScheduleUpdateCount: 0, + owningRendererScheduleUpdateCount: 2, + rapidCycleScheduleUpdateCount: 1, + staleScheduleUpdateCount: 0, + }); + }); + test("should handle multiple prompt mode cycles correctly", async ({ reactGrab }) => { await reactGrab.registerCommentAction(); diff --git a/packages/react-grab/src/utils/freeze-updates.ts b/packages/react-grab/src/utils/freeze-updates.ts index 4c3eea635..e69d9e4e5 100644 --- a/packages/react-grab/src/utils/freeze-updates.ts +++ b/packages/react-grab/src/utils/freeze-updates.ts @@ -9,6 +9,7 @@ import { _fiberRoots, getFiberFromHostInstance, getRDTHook, + instrument, isCompositeFiber, type Fiber, type ReactRenderer, @@ -61,6 +62,7 @@ interface PausedContextState { let isUpdatesPaused = false; let freezeOwnerCount = 0; +let freezeSessionId = 0; const getOrCache = (cache: WeakMap, key: K, create: () => V): V => { const cached = cache.get(key); @@ -91,9 +93,24 @@ const pausedContextStates = new WeakMap() const renderersWithPatchedDispatcher = new WeakSet(); const typedFiberRoots = _fiberRoots as Set; const pausedFiberRoots = new Set(); +const fiberRootRenderers = new WeakMap(); -const isDomRenderer = (renderer: ReactRenderer): boolean => - renderer.rendererPackageName.startsWith("react-dom"); +instrument({ + name: "react-grab-freeze-updates", + onCommitFiberRoot: (rendererId, fiberRoot) => { + const renderer = getRDTHook().renderers.get(rendererId); + if (renderer) fiberRootRenderers.set(fiberRoot, renderer); + }, +}); + +const isDomRenderer = (renderer: ReactRenderer): boolean => { + try { + const packageName = renderer.rendererPackageName; + return typeof packageName === "string" && packageName.startsWith("react-dom"); + } catch { + return false; + } +}; const getFiberRoot = (fiber: Fiber): FiberRootLike | null => { let current: Fiber | null = fiber; @@ -103,6 +120,58 @@ const getFiberRoot = (fiber: Fiber): FiberRootLike | null => { return (current.stateNode ?? null) as FiberRootLike | null; }; +const findHostInstance = (fiberRoot: FiberRootLike): object | null => { + const root = fiberRoot.current; + let fiber = root; + while (fiber) { + const stateNode = fiber.stateNode; + if ( + stateNode && + typeof stateNode === "object" && + typeof Reflect.get(stateNode, "nodeType") === "number" + ) { + return stateNode; + } + if (fiber.child) { + fiber = fiber.child; + continue; + } + while (fiber !== root && !fiber.sibling) { + fiber = fiber.return; + if (!fiber) return null; + } + if (fiber === root) return null; + fiber = fiber.sibling; + } + return null; +}; + +const resolveFiberRootRenderer = (fiberRoot: FiberRootLike): ReactRenderer | null => { + const renderer = fiberRootRenderers.get(fiberRoot); + if (renderer) return isDomRenderer(renderer) ? renderer : null; + + const domRenderers = Array.from(getRDTHook().renderers.values()).filter(isDomRenderer); + if (domRenderers.length === 1) { + const domRenderer = domRenderers[0]; + if (domRenderer) fiberRootRenderers.set(fiberRoot, domRenderer); + return domRenderer ?? null; + } + + const hostInstance = findHostInstance(fiberRoot); + if (!hostInstance) return null; + + for (const domRenderer of domRenderers) { + try { + const hostFiber = domRenderer.findFiberByHostInstance?.(hostInstance); + if (hostFiber && getFiberRoot(hostFiber) === fiberRoot) { + fiberRootRenderers.set(fiberRoot, domRenderer); + return domRenderer; + } + } catch {} + } + return null; +}; + const isDomFiberRoot = (fiberRoot: FiberRootLike): boolean => { const stateNode = fiberRoot.current?.stateNode; if (!stateNode || typeof stateNode !== "object") return false; @@ -531,21 +600,22 @@ const installDispatcherPatching = (renderer: ReactRenderer): void => { }); }; -const scheduleReactUpdate = (fiberRoots: Set): void => { +const scheduleReactUpdate = ( + fiberRoots: Set, + scheduledFreezeSessionId: number, +): void => { queueMicrotask(() => { + if (isUpdatesPaused || freezeSessionId !== scheduledFreezeSessionId) return; try { - for (const renderer of getRDTHook().renderers.values()) { - if (!isDomRenderer(renderer)) continue; - if (typeof renderer.scheduleUpdate !== "function") continue; - for (const fiberRoot of fiberRoots) { - if (fiberRoot.current) { - try { - renderer.scheduleUpdate(fiberRoot.current); - } catch (error) { - reportRecoverableError( - new RecoverableError("scheduleUpdate failed during unfreeze", error), - ); - } + for (const fiberRoot of fiberRoots) { + const renderer = resolveFiberRootRenderer(fiberRoot); + if (fiberRoot.current && renderer?.scheduleUpdate) { + try { + renderer.scheduleUpdate(fiberRoot.current); + } catch (error) { + reportRecoverableError( + new RecoverableError("scheduleUpdate failed during unfreeze", error), + ); } } } @@ -581,6 +651,7 @@ const clearPendingUpdates = (): void => { }; const resumeUpdates = (): void => { + const resumedFreezeSessionId = freezeSessionId; const fiberRootsToResume = new Set(pausedFiberRoots); try { for (const fiberRoot of collectFiberRoots()) { @@ -613,8 +684,8 @@ const resumeUpdates = (): void => { invokeCallbacks(storeCallbacksToInvoke); invokeCallbacks(transitionCallbacksToInvoke); invokeCallbacks(stateUpdatesToInvoke); - if (!isUpdatesPaused) { - scheduleReactUpdate(fiberRootsToResume); + if (!isUpdatesPaused && freezeSessionId === resumedFreezeSessionId) { + scheduleReactUpdate(fiberRootsToResume, resumedFreezeSessionId); } }; @@ -629,6 +700,7 @@ export const freezeUpdatesOrThrow = (): (() => void) => { if (isFirstFreezeOwner) { try { initializeFreezeSupport(); + freezeSessionId += 1; isUpdatesPaused = true; const fiberRoots = collectFiberRoots();