feat(react/runtime): add context tracking for worklet execution

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

@@ -4,6 +4,7 @@
 import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { initApiEnv } from '../src/api/lynxApi';
+import { currentCtx } from '../src/ctxTrace';
 import { updateWorkletRefInitValueChanges } from '../src/workletRef';
 import { initWorklet } from '../src/workletRuntime';
 
@@ -154,6 +155,21 @@ describe('Worklet', () => {
     expect(fn).toHaveBeenLastCalledWith(4);
   });
 
+  it('should track current worklet ctx', async () => {
+    initWorklet();
+
+    const fn = vi.fn(() => {
+      expect(currentCtx.ctx._wkltId).toEqual('1');
+    });
+    registerWorklet('main-thread', '1', fn);
+    let worklet = {
+      _wkltId: '1',
+    };
+    runWorklet(worklet);
+    expect(fn).toBeCalled();
+    expect(currentCtx).toBeUndefined();
+  });
+
   it('value of a workletRef should be preserved between calls', async () => {
     initWorklet();
 

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

@@ -0,0 +1,21 @@
+// 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 type { ClosureValueType, Worklet } from './bindings/types.js';
+
+interface CtxTrace {
+  ctx: Worklet;
+}
+
+export let currentCtx: CtxTrace | undefined;
+
+export function traceCtxCall(ctx: Worklet, _params: ClosureValueType[]): void {
+  currentCtx = {
+    ctx,
+  };
+}
+
+export function clearCurrentCtx(): void {
+  currentCtx = undefined;
+}

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

@@ -3,8 +3,9 @@
 // LICENSE file in the root directory of this source tree.
 import { Element } from './api/element.js';
 import type { ClosureValueType, Worklet, WorkletRefImpl } from './bindings/types.js';
+import { clearCurrentCtx, traceCtxCall } from './ctxTrace.js';
 import { delayExecUntilJsReady, initEventDelay } from './delayWorkletEvent.js';
-import { isRunOnBackgroundEnabled, JsFunctionLifecycleManager } from './jsFunctionLifecycle.js';
+import { JsFunctionLifecycleManager, isRunOnBackgroundEnabled } from './jsFunctionLifecycle.js';
 import { profile } from './utils/profile.js';
 import { getFromWorkletRefMap, initWorkletRef } from './workletRef.js';
 
@@ -54,20 +55,20 @@ function runWorklet(ctx: Worklet, params: ClosureValueType[]): unknown {
 }
 
 function runWorkletImpl(ctx: Worklet, params: ClosureValueType[]): unknown {
-  const worklet: Function = profile(
-    'transformWorkletCtx ' + ctx._wkltId,
-    () => transformWorklet(ctx, true),
-  );
-  const params_: ClosureValueType[] = profile(
-    'transformWorkletParams',
-    () => transformWorklet(params || [], false),
-  );
-
-  let result;
-  profile('runWorklet', () => {
-    result = worklet(...params_);
-  });
-  return result;
+  traceCtxCall(ctx, params);
+  try {
+    const worklet: (...args: any[]) => unknown = profile(
+      'transformWorkletCtx ' + ctx._wkltId,
+      () => transformWorklet(ctx, true),
+    );
+    const params_: ClosureValueType[] = profile(
+      'transformWorkletParams',
+      () => transformWorklet(params || [], false),
+    );
+    return profile('runWorklet', () => worklet(...params_));
+  } finally {
+    clearCurrentCtx();
+  }
 }
 
 function validateWorklet(ctx: unknown): ctx is Worklet {