diff --git a/packages/react/worklet-runtime/__test__/workletRuntime.test.js b/packages/react/worklet-runtime/__test__/workletRuntime.test.js index 9afb532df3..e294e92b23 100644 --- a/packages/react/worklet-runtime/__test__/workletRuntime.test.js +++ b/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(); diff --git a/packages/react/worklet-runtime/src/ctxTrace.ts b/packages/react/worklet-runtime/src/ctxTrace.ts new file mode 100644 index 0000000000..c7c4843282 --- /dev/null +++ b/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; +} diff --git a/packages/react/worklet-runtime/src/workletRuntime.ts b/packages/react/worklet-runtime/src/workletRuntime.ts index 6ecaf15fb2..5d2af97178 100644 --- a/packages/react/worklet-runtime/src/workletRuntime.ts +++ b/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 {