fix(react/runtime): delay applying mt:ref until hydration has finished

.changeset/red-meals-battle.md

@@ -0,0 +1,5 @@
+---
+"@lynx-js/react": patch
+---
+
+Fix a problem causing `MainThreadRef`s to not be updated correctly during hydration when they are set to `main-thread:ref`s.

packages/react/runtime/__test__/snapshot/workletRef.test.jsx

@@ -30,6 +30,11 @@ beforeAll(() => {
     _runOnBackgroundDelayImpl: {
       runDelayedBackgroundFunctions: vi.fn(),
     },
+    _eventDelayImpl: {
+      runDelayedWorklet: vi.fn(),
+      clearDelayedWorklets: vi.fn(),
+    },
+    _hydrateCtx: vi.fn(),
   };
   globalThis.runWorklet = vi.fn();
 });
@@ -334,25 +339,25 @@ describe('WorkletRef', () => {
           ],
           [
             {
-              "_execId": 2,
               "_unmount": undefined,
               "_wkltId": 233,
             },
             [
-              {
-                "elementRefptr": <view
-                  has-react-ref={true}
-                />,
-              },
+              null,
             ],
           ],
           [
             {
+              "_execId": 2,
               "_unmount": undefined,
               "_wkltId": 233,
             },
             [
-              null,
+              {
+                "elementRefptr": <view
+                  has-react-ref={true}
+                />,
+              },
             ],
           ],
           [
@@ -371,7 +376,7 @@ describe('WorkletRef', () => {
           ],
         ]
       `);
-      globalThis.runWorklet.mock.calls[1][0]._unmount = cleanup;
+      globalThis.runWorklet.mock.calls[2][0]._unmount = cleanup;
     }
 
     // update

packages/react/runtime/__test__/worklet/workletRef.test.jsx

@@ -28,6 +28,9 @@ beforeAll(() => {
     _runOnBackgroundDelayImpl: {
       runDelayedBackgroundFunctions: vi.fn(),
     },
+    _eventDelayImpl: {
+      clearDelayedWorklets: vi.fn(),
+    },
   };
 });
 

packages/react/runtime/src/lifecycle/patch/updateMainThread.ts

@@ -6,12 +6,13 @@ import { updateWorkletRefInitValueChanges } from '@lynx-js/react/worklet-runtime
 
 import { LifecycleConstant } from '../../lifecycleConstant.js';
 import { __pendingListUpdates } from '../../list.js';
-import { markTiming, PerformanceTimingKeys, setPipeline } from '../../lynx/performance.js';
+import { PerformanceTimingKeys, markTiming, setPipeline } from '../../lynx/performance.js';
 import { __page } from '../../snapshot.js';
 import { getReloadVersion } from '../pass.js';
 import type { PatchList, PatchOptions } from './commit.js';
 import { setMainThreadHydrationFinished } from './isMainThreadHydrationFinished.js';
 import { snapshotPatchApply } from './snapshotPatchApply.js';
+import { applyRefQueue } from '../../snapshot/workletRef.js';
 
 function updateMainThread(
   { data, patchOptions }: {
@@ -46,6 +47,7 @@ function updateMainThread(
   if (patchOptions.isHydration) {
     setMainThreadHydrationFinished(true);
   }
+  applyRefQueue();
   if (patchOptions.pipelineOptions) {
     flushOptions.pipelineOptions = patchOptions.pipelineOptions;
   }

packages/react/runtime/src/lifecycle/reload.ts

@@ -13,16 +13,17 @@ import { hydrate } from '../hydrate.js';
 import { LifecycleConstant } from '../lifecycleConstant.js';
 import { __pendingListUpdates } from '../list.js';
 import { __root, setRoot } from '../root.js';
+import { destroyBackground } from './destroy.js';
+import { applyRefQueue } from '../snapshot/workletRef.js';
 import { SnapshotInstance, __page, snapshotInstanceManager } from '../snapshot.js';
 import { isEmptyObject } from '../utils.js';
-import { destroyBackground } from './destroy.js';
 import { destroyWorklet } from '../worklet/destroy.js';
 import { clearJSReadyEventIdSwap, isJSReady } from './event/jsReady.js';
 import { increaseReloadVersion } from './pass.js';
+import { setMainThreadHydrationFinished } from './patch/isMainThreadHydrationFinished.js';
 import { deinitGlobalSnapshotPatch } from './patch/snapshotPatch.js';
 import { shouldDelayUiOps } from './ref/delay.js';
 import { renderMainThread } from './render.js';
-import { setMainThreadHydrationFinished } from './patch/isMainThreadHydrationFinished.js';
 
 function reloadMainThread(data: any, options: UpdatePageOption): void {
   if (__PROFILE__) {
@@ -51,6 +52,7 @@ function reloadMainThread(data: any, options: UpdatePageOption): void {
 
   // always call this before `__FlushElementTree`
   __pendingListUpdates.flush();
+  applyRefQueue();
 
   if (isJSReady) {
     __OnLifecycleEvent([

packages/react/runtime/src/list.ts

@@ -2,6 +2,7 @@
 // Licensed under the Apache License Version 2.0 that can be found in the
 // LICENSE file in the root directory of this source tree.
 import { hydrate } from './hydrate.js';
+import { applyRefQueue } from './snapshot/workletRef.js';
 import type { SnapshotInstance } from './snapshot.js';
 
 export interface ListUpdateInfo {
@@ -311,6 +312,7 @@ export function componentAtIndexFactory(
       hydrate(oldCtx, childCtx);
       oldCtx.unRenderElements();
       const root = childCtx.__element_root!;
+      applyRefQueue();
       if (!enableBatchRender) {
         const flushOptions: FlushOptions = {
           triggerLayout: true,
@@ -345,6 +347,7 @@ export function componentAtIndexFactory(
     const root = childCtx.__element_root!;
     __AppendElement(list, root);
     const sign = __GetElementUniqueID(root);
+    applyRefQueue();
     if (!enableBatchRender) {
       __FlushElementTree(root, {
         triggerLayout: true,

packages/react/runtime/src/lynx/calledByNative.ts

@@ -12,6 +12,7 @@ import { __root, setRoot } from '../root.js';
 import { SnapshotInstance, __page, setupPage } from '../snapshot.js';
 import { isEmptyObject } from '../utils.js';
 import { PerformanceTimingKeys, markTiming, setPipeline } from './performance.js';
+import { applyRefQueue } from '../snapshot/workletRef.js';
 
 function ssrEncode() {
   const { __opcodes } = __root;
@@ -87,6 +88,7 @@ function renderPage(data: any): void {
   // always call this before `__FlushElementTree`
   // (There is an implicit `__FlushElementTree` in `renderPage`)
   __pendingListUpdates.flush();
+  applyRefQueue();
 
   if (__FIRST_SCREEN_SYNC_TIMING__ === 'immediately') {
     jsReady();
@@ -128,6 +130,7 @@ function updatePage(data: any, options?: UpdatePageOption): void {
 
       // always call this before `__FlushElementTree`
       __pendingListUpdates.flush();
+      applyRefQueue();
     }
     markTiming(PerformanceTimingKeys.updateDiffVdomEnd);
   }

packages/react/runtime/src/snapshot/workletRef.ts

@@ -1,12 +1,34 @@
 // Copyright 2024 The Lynx Authors. All rights reserved.
 // Licensed under the Apache License Version 2.0 that can be found in the
 // LICENSE file in the root directory of this source tree.
-import { runWorkletCtx, updateWorkletRef as update } from '@lynx-js/react/worklet-runtime/bindings';
+
+import { onWorkletCtxUpdate, runWorkletCtx, updateWorkletRef as update } from '@lynx-js/react/worklet-runtime/bindings';
 import type { Element, Worklet, WorkletRefImpl } from '@lynx-js/react/worklet-runtime/bindings';
 
+import { isMainThreadHydrationFinished } from '../lifecycle/patch/isMainThreadHydrationFinished.js';
 import type { SnapshotInstance } from '../snapshot.js';
 
-function workletUnRef(value: Worklet | WorkletRefImpl<Element>): void {
+let mtRefQueue: (WorkletRefImpl<Element> | Worklet | Element)[] = [];
+
+export function applyRefQueue(): void {
+  const queue = mtRefQueue;
+  mtRefQueue = [];
+  for (let i = 0; i < queue.length; i += 2) {
+    const worklet = queue[i] as Worklet | WorkletRefImpl<Element>;
+    const element = queue[i + 1] as Element;
+    if ('_wvid' in worklet) {
+      update(worklet as WorkletRefImpl<Element>, element);
+    } else if ('_wkltId' in worklet) {
+      worklet._unmount = runWorkletCtx(worklet, [{ elementRefptr: element }]) as () => void;
+    }
+  }
+}
+
+function addToRefQueue(worklet: Worklet | WorkletRefImpl<Element>, element: Element): void {
+  mtRefQueue.push(worklet, element);
+}
+
+export function workletUnRef(value: Worklet | WorkletRefImpl<Element>): void {
   if ('_wvid' in value) {
     update(value as WorkletRefImpl<Element>, null);
   } else if ('_wkltId' in value) {
@@ -18,7 +40,7 @@ function workletUnRef(value: Worklet | WorkletRefImpl<Element>): void {
   }
 }
 
-function updateWorkletRef(
+export function updateWorkletRef(
   snapshot: SnapshotInstance,
   expIndex: number,
   oldValue: WorkletRefImpl<Element> | Worklet | undefined,
@@ -38,11 +60,17 @@ function updateWorkletRef(
   if (value === null || value === undefined) {
     // do nothing
   } else if (value._wvid) {
-    update(value as WorkletRefImpl<Element>, snapshot.__elements[elementIndex]!);
+    const element = snapshot.__elements[elementIndex]! as Element;
+    addToRefQueue(value as Worklet, element);
   } else if ((value as Worklet)._wkltId) {
-    (value as Worklet)._unmount = runWorkletCtx(value as Worklet, [{
-      elementRefptr: (snapshot.__elements[elementIndex]!) as any,
-    }]) as () => void;
+    const element = snapshot.__elements[elementIndex]! as Element;
+    onWorkletCtxUpdate(
+      value as Worklet,
+      oldValue as Worklet | undefined,
+      !isMainThreadHydrationFinished,
+      element,
+    );
+    addToRefQueue(value as Worklet, element);
     /* v8 ignore next 3 */
   } else if (value._type === '__LEPUS__' || (value as Worklet)._lepusWorkletHash) {
     // for pre-0.99 compatibility
@@ -58,5 +86,3 @@ function updateWorkletRef(
   // Add an arbitrary attribute to avoid this element being layout-only
   __SetAttribute(snapshot.__elements[elementIndex]!, 'has-react-ref', true);
 }
-
-export { updateWorkletRef, workletUnRef };

packages/react/worklet-runtime/__test__/workletRef.test.js

@@ -134,11 +134,11 @@ describe('WorkletRef', () => {
       [4, 'background-thread-init-4'],
       [5, 'background-thread-init-5'],
     ]);
+    globalThis.lynxWorkletImpl._hydrateCtx(worklet, firstScreenWorklet);
     globalThis.lynxWorkletImpl._refImpl.updateWorkletRef({
       _wvid: 5,
       _initValue: 'background-thread-init-5',
     }, 'background-thread-element-5');
-    globalThis.lynxWorkletImpl._hydrateCtx(worklet, firstScreenWorklet);
     expect(getFromWorkletRefMap({ _wvid: 1 }).current).toBe('main-thread-set-1');
     expect(getFromWorkletRefMap({ _wvid: 2 }).current).toBe('main-thread-init-2');
     expect(getFromWorkletRefMap({ _wvid: 3 }).current).toBe('main-thread-init-3');

packages/react/worklet-runtime/src/bindings/observers.ts

@@ -11,9 +11,8 @@ import type { Worklet } from './types.js';
  * @param oldWorklet - The old worklet context
  * @param isFirstScreen - Whether it is before the hydration is finished
  * @param element - The element
- * @internal
  */
-function onWorkletCtxUpdate(
+export function onWorkletCtxUpdate(
   worklet: Worklet,
   oldWorklet: Worklet | null | undefined,
   isFirstScreen: boolean,
@@ -24,17 +23,17 @@ function onWorkletCtxUpdate(
     globalThis.lynxWorkletImpl?._hydrateCtx(worklet, oldWorklet);
   }
   // For old version dynamic component compatibility.
-  globalThis.lynxWorkletImpl?._eventDelayImpl.runDelayedWorklet(worklet, element);
+  if (isFirstScreen) {
+    globalThis.lynxWorkletImpl?._eventDelayImpl.runDelayedWorklet(worklet, element);
+  }
 }
 
 /**
  * This must be called when the hydration is finished.
- *
- * @internal
  */
-function onHydrationFinished(): void {
+export function onHydrationFinished(): void {
   globalThis.lynxWorkletImpl?._runOnBackgroundDelayImpl.runDelayedBackgroundFunctions();
   globalThis.lynxWorkletImpl?._refImpl.clearFirstScreenWorkletRefMap();
+  // For old version dynamic component compatibility.
+  globalThis.lynxWorkletImpl?._eventDelayImpl.clearDelayedWorklets();
 }
-
-export { onWorkletCtxUpdate, onHydrationFinished };

packages/react/worklet-runtime/src/hydrate.ts

@@ -2,7 +2,6 @@
 // Licensed under the Apache License Version 2.0 that can be found in the
 // LICENSE file in the root directory of this source tree.
 
-import { Element } from './api/element.js';
 import type { ClosureValueType, JsFnHandle, Worklet, WorkletRefId, WorkletRefImpl } from './bindings/index.js';
 import { profile } from './utils/profile.js';
 
@@ -64,10 +63,6 @@ function hydrateMainThreadRef(refId: WorkletRefId, value: WorkletRefImpl<unknown
     return;
   }
   const ref = lynxWorkletImpl!._refImpl._workletRefMap[refId]!;
-  if (ref.current instanceof Element) {
-    // Modified by `main-thread:ref`
-    return;
-  }
   ref.current = value.current;
 }