From c05b8b5054730975d6f1faece210b6725219b185 Mon Sep 17 00:00:00 2001 From: Sebastian Markbage Date: Tue, 9 Apr 2024 21:41:25 -0400 Subject: [PATCH 1/3] Add component storage async context This will store owner on a per component basis. --- .../src/forks/ReactFlightServerConfig.custom.js | 5 +++++ .../forks/ReactFlightServerConfig.dom-browser-esm.js | 11 +++++++---- .../ReactFlightServerConfig.dom-browser-turbopack.js | 5 +++++ .../src/forks/ReactFlightServerConfig.dom-browser.js | 5 +++++ .../src/forks/ReactFlightServerConfig.dom-bun.js | 5 +++++ .../ReactFlightServerConfig.dom-edge-turbopack.js | 6 ++++++ .../src/forks/ReactFlightServerConfig.dom-edge.js | 7 +++++++ .../src/forks/ReactFlightServerConfig.dom-legacy.js | 5 +++++ .../src/forks/ReactFlightServerConfig.dom-node-esm.js | 6 ++++++ .../ReactFlightServerConfig.dom-node-turbopack.js | 5 +++++ .../src/forks/ReactFlightServerConfig.dom-node.js | 5 +++++ 11 files changed, 61 insertions(+), 4 deletions(-) diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.custom.js b/packages/react-server/src/forks/ReactFlightServerConfig.custom.js index b78f9e3816f37..5afdf8c29c888 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.custom.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.custom.js @@ -8,6 +8,7 @@ */ import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from '../ReactFlightServerConfigBundlerCustom'; @@ -23,6 +24,10 @@ export const isPrimaryRenderer = false; export const supportsRequestStorage = false; export const requestStorage: AsyncLocalStorage = (null: any); +export const supportsComponentStorage = false; +export const componentStorage: AsyncLocalStorage = + (null: any); + export function createHints(): any { return null; } diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-esm.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-esm.js index 929c2707c38d9..3b6251d37b4b3 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-esm.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-esm.js @@ -6,15 +6,18 @@ * * @flow */ -import {AsyncLocalStorage} from 'async_hooks'; import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-esm/src/ReactFlightServerConfigESMBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; -export const supportsRequestStorage = true; -export const requestStorage: AsyncLocalStorage = - new AsyncLocalStorage(); +export const supportsRequestStorage = false; +export const requestStorage: AsyncLocalStorage = (null: any); + +export const supportsComponentStorage = false; +export const componentStorage: AsyncLocalStorage = + (null: any); export * from '../ReactFlightServerConfigDebugNoop'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-turbopack.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-turbopack.js index 205a7add5fdfc..9b9440a3dae80 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-turbopack.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-turbopack.js @@ -8,6 +8,7 @@ */ import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; export const supportsRequestStorage = false; export const requestStorage: AsyncLocalStorage = (null: any); +export const supportsComponentStorage = false; +export const componentStorage: AsyncLocalStorage = + (null: any); + export * from '../ReactFlightServerConfigDebugNoop'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser.js index c5f4407796161..4dbfae1ce3ecf 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-browser.js @@ -8,6 +8,7 @@ */ import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; export const supportsRequestStorage = false; export const requestStorage: AsyncLocalStorage = (null: any); +export const supportsComponentStorage = false; +export const componentStorage: AsyncLocalStorage = + (null: any); + export * from '../ReactFlightServerConfigDebugNoop'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-bun.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-bun.js index ad1743a43d18a..bfe794c9a2c55 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-bun.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-bun.js @@ -8,6 +8,7 @@ */ import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from '../ReactFlightServerConfigBundlerCustom'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; export const supportsRequestStorage = false; export const requestStorage: AsyncLocalStorage = (null: any); +export const supportsComponentStorage = false; +export const componentStorage: AsyncLocalStorage = + (null: any); + export * from '../ReactFlightServerConfigDebugNoop'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge-turbopack.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge-turbopack.js index 8b33307667574..56d00364ec247 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge-turbopack.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge-turbopack.js @@ -7,6 +7,7 @@ * @flow */ import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -16,6 +17,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function'; export const requestStorage: AsyncLocalStorage = supportsRequestStorage ? new AsyncLocalStorage() : (null: any); +export const supportsComponentStorage: boolean = + __DEV__ && supportsRequestStorage; +export const componentStorage: AsyncLocalStorage = + supportsComponentStorage ? new AsyncLocalStorage() : (null: any); + // We use the Node version but get access to async_hooks from a global. import type {HookCallbacks, AsyncHook} from 'async_hooks'; export const createAsyncHook: HookCallbacks => AsyncHook = diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge.js index e62460390d19e..322f3a10989d4 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-edge.js @@ -6,7 +6,9 @@ * * @flow */ + import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -16,6 +18,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function'; export const requestStorage: AsyncLocalStorage = supportsRequestStorage ? new AsyncLocalStorage() : (null: any); +export const supportsComponentStorage: boolean = + __DEV__ && supportsRequestStorage; +export const componentStorage: AsyncLocalStorage = + supportsComponentStorage ? new AsyncLocalStorage() : (null: any); + // We use the Node version but get access to async_hooks from a global. import type {HookCallbacks, AsyncHook} from 'async_hooks'; export const createAsyncHook: HookCallbacks => AsyncHook = diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-legacy.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-legacy.js index ad1743a43d18a..bfe794c9a2c55 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-legacy.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-legacy.js @@ -8,6 +8,7 @@ */ import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from '../ReactFlightServerConfigBundlerCustom'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; export const supportsRequestStorage = false; export const requestStorage: AsyncLocalStorage = (null: any); +export const supportsComponentStorage = false; +export const componentStorage: AsyncLocalStorage = + (null: any); + export * from '../ReactFlightServerConfigDebugNoop'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-esm.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-esm.js index 528c3cdb3b23b..b44d6fe2457b8 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-esm.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-esm.js @@ -6,9 +6,11 @@ * * @flow */ + import {AsyncLocalStorage} from 'async_hooks'; import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-esm/src/ReactFlightServerConfigESMBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -17,5 +19,9 @@ export const supportsRequestStorage = true; export const requestStorage: AsyncLocalStorage = new AsyncLocalStorage(); +export const supportsComponentStorage = __DEV__; +export const componentStorage: AsyncLocalStorage = + supportsComponentStorage ? new AsyncLocalStorage() : (null: any); + export {createHook as createAsyncHook, executionAsyncId} from 'async_hooks'; export * from '../ReactFlightServerConfigDebugNode'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-turbopack.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-turbopack.js index a19e29222403c..eb566f3e2e43c 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-turbopack.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-node-turbopack.js @@ -10,6 +10,7 @@ import {AsyncLocalStorage} from 'async_hooks'; import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -18,5 +19,9 @@ export const supportsRequestStorage = true; export const requestStorage: AsyncLocalStorage = new AsyncLocalStorage(); +export const supportsComponentStorage = __DEV__; +export const componentStorage: AsyncLocalStorage = + supportsComponentStorage ? new AsyncLocalStorage() : (null: any); + export {createHook as createAsyncHook, executionAsyncId} from 'async_hooks'; export * from '../ReactFlightServerConfigDebugNode'; diff --git a/packages/react-server/src/forks/ReactFlightServerConfig.dom-node.js b/packages/react-server/src/forks/ReactFlightServerConfig.dom-node.js index b0da1d9926e68..40c58ea2dde12 100644 --- a/packages/react-server/src/forks/ReactFlightServerConfig.dom-node.js +++ b/packages/react-server/src/forks/ReactFlightServerConfig.dom-node.js @@ -10,6 +10,7 @@ import {AsyncLocalStorage} from 'async_hooks'; import type {Request} from 'react-server/src/ReactFlightServer'; +import type {ReactComponentInfo} from 'shared/ReactTypes'; export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler'; export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM'; @@ -18,5 +19,9 @@ export const supportsRequestStorage = true; export const requestStorage: AsyncLocalStorage = new AsyncLocalStorage(); +export const supportsComponentStorage = __DEV__; +export const componentStorage: AsyncLocalStorage = + supportsComponentStorage ? new AsyncLocalStorage() : (null: any); + export {createHook as createAsyncHook, executionAsyncId} from 'async_hooks'; export * from '../ReactFlightServerConfigDebugNode'; From df153c96510de9fd039c8d5fb5a63e143b6d315b Mon Sep 17 00:00:00 2001 From: Sebastian Markbage Date: Tue, 9 Apr 2024 21:50:55 -0400 Subject: [PATCH 2/3] Split out current owner for parity with Fiber --- packages/react-server/src/ReactFlightServer.js | 8 +++----- .../src/flight/ReactFlightAsyncDispatcher.js | 8 ++------ .../src/flight/ReactFlightCurrentOwner.js | 16 ++++++++++++++++ 3 files changed, 21 insertions(+), 11 deletions(-) create mode 100644 packages/react-server/src/flight/ReactFlightCurrentOwner.js diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js index 5c66d56bf3075..78d502c0d149d 100644 --- a/packages/react-server/src/ReactFlightServer.js +++ b/packages/react-server/src/ReactFlightServer.js @@ -89,11 +89,9 @@ import { getThenableStateAfterSuspending, resetHooksForRequest, } from './ReactFlightHooks'; -import { - DefaultAsyncDispatcher, - currentOwner, - setCurrentOwner, -} from './flight/ReactFlightAsyncDispatcher'; +import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher'; + +import {currentOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner'; import { getIteratorFn, diff --git a/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js b/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js index 1cdd67a82153b..3ccd00b466459 100644 --- a/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js +++ b/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js @@ -15,6 +15,8 @@ import {resolveRequest, getCache} from '../ReactFlightServer'; import {disableStringRefs} from 'shared/ReactFeatureFlags'; +import {currentOwner} from './ReactFlightCurrentOwner'; + function resolveCache(): Map { const request = resolveRequest(); if (request) { @@ -36,8 +38,6 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({ }, }: any); -export let currentOwner: ReactComponentInfo | null = null; - if (__DEV__) { DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => { return currentOwner; @@ -48,7 +48,3 @@ if (__DEV__) { return null; }; } - -export function setCurrentOwner(componentInfo: null | ReactComponentInfo) { - currentOwner = componentInfo; -} diff --git a/packages/react-server/src/flight/ReactFlightCurrentOwner.js b/packages/react-server/src/flight/ReactFlightCurrentOwner.js new file mode 100644 index 0000000000000..6b5d2dd46a8b4 --- /dev/null +++ b/packages/react-server/src/flight/ReactFlightCurrentOwner.js @@ -0,0 +1,16 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + * + * @flow + */ + +import type {ReactComponentInfo} from 'shared/ReactTypes'; + +export let currentOwner: ReactComponentInfo | null = null; + +export function setCurrentOwner(componentInfo: null | ReactComponentInfo) { + currentOwner = componentInfo; +} From a2d2524932459b65a226f50808b7028b26e9f26c Mon Sep 17 00:00:00 2001 From: Sebastian Markbage Date: Tue, 9 Apr 2024 22:50:39 -0400 Subject: [PATCH 3/3] Track owner on componentStorage --- .../src/__tests__/ReactFlightDOMEdge-test.js | 71 +++++++++++++++++++ .../react-server/src/ReactFlightServer.js | 33 ++++++--- .../src/flight/ReactFlightAsyncDispatcher.js | 6 +- .../src/flight/ReactFlightCurrentOwner.js | 16 ++++- 4 files changed, 110 insertions(+), 16 deletions(-) diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js index 6e2e02047bbe6..6bbbf9d82f69b 100644 --- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js +++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js @@ -21,6 +21,8 @@ if (typeof Blob === 'undefined') { if (typeof File === 'undefined') { global.File = require('buffer').File; } +// Patch for Edge environments for global scope +global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage; // Don't wait before processing work on the server. // TODO: we can replace this with FlightServer.act(). @@ -32,6 +34,7 @@ let webpackMap; let webpackModules; let webpackModuleLoading; let React; +let ReactServer; let ReactDOMServer; let ReactServerDOMServer; let ReactServerDOMClient; @@ -55,6 +58,7 @@ describe('ReactFlightDOMEdge', () => { webpackModules = WebpackMock.webpackModules; webpackModuleLoading = WebpackMock.moduleLoading; + ReactServer = require('react'); ReactServerDOMServer = require('react-server-dom-webpack/server'); jest.resetModules(); @@ -692,4 +696,71 @@ describe('ReactFlightDOMEdge', () => { ), ); }); + + it('supports async server component debug info as the element owner in DEV', async () => { + function Container({children}) { + return children; + } + + const promise = Promise.resolve(true); + async function Greeting({firstName}) { + // We can't use JSX here because it'll use the Client React. + const child = ReactServer.createElement( + 'span', + null, + 'Hello, ' + firstName, + ); + // Yield the synchronous pass + await promise; + // We should still be able to track owner using AsyncLocalStorage. + return ReactServer.createElement(Container, null, child); + } + + const model = { + greeting: ReactServer.createElement(Greeting, {firstName: 'Seb'}), + }; + + const stream = ReactServerDOMServer.renderToReadableStream( + model, + webpackMap, + ); + + const rootModel = await ReactServerDOMClient.createFromReadableStream( + stream, + { + ssrManifest: { + moduleMap: null, + moduleLoading: null, + }, + }, + ); + + const ssrStream = await ReactDOMServer.renderToReadableStream( + rootModel.greeting, + ); + const result = await readResult(ssrStream); + expect(result).toEqual('Hello, Seb'); + + // Resolve the React Lazy wrapper which must have resolved by now. + const lazyWrapper = rootModel.greeting; + const greeting = lazyWrapper._init(lazyWrapper._payload); + + // We've rendered down to the span. + expect(greeting.type).toBe('span'); + if (__DEV__) { + const greetInfo = {name: 'Greeting', env: 'Server', owner: null}; + expect(lazyWrapper._debugInfo).toEqual([ + greetInfo, + {name: 'Container', env: 'Server', owner: greetInfo}, + ]); + // The owner that created the span was the outer server component. + // We expect the debug info to be referentially equal to the owner. + expect(greeting._owner).toBe(lazyWrapper._debugInfo[0]); + } else { + expect(lazyWrapper._debugInfo).toBe(undefined); + expect(greeting._owner).toBe( + gate(flags => flags.disableStringRefs) ? undefined : null, + ); + } + }); }); diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js index 78d502c0d149d..13dd33528a86b 100644 --- a/packages/react-server/src/ReactFlightServer.js +++ b/packages/react-server/src/ReactFlightServer.js @@ -73,6 +73,8 @@ import { isServerReference, supportsRequestStorage, requestStorage, + supportsComponentStorage, + componentStorage, createHints, initAsyncDebugInfo, } from './ReactFlightServerConfig'; @@ -91,7 +93,7 @@ import { } from './ReactFlightHooks'; import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher'; -import {currentOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner'; +import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner'; import { getIteratorFn, @@ -160,7 +162,7 @@ function patchConsole(consoleInst: typeof console, methodName: string) { // We don't currently use this id for anything but we emit it so that we can later // refer to previous logs in debug info to associate them with a component. const id = request.nextChunkId++; - const owner: null | ReactComponentInfo = currentOwner; + const owner: null | ReactComponentInfo = resolveOwner(); emitConsoleChunk(request, id, methodName, owner, stack, arguments); } // $FlowFixMe[prop-missing] @@ -822,7 +824,11 @@ function renderFunctionComponent( const prevThenableState = task.thenableState; task.thenableState = null; - let componentDebugInfo: null | ReactComponentInfo = null; + // The secondArg is always undefined in Server Components since refs error early. + const secondArg = undefined; + let result; + + let componentDebugInfo: ReactComponentInfo; if (__DEV__) { if (debugID === null) { // We don't have a chunk to assign debug info. We need to outline this @@ -851,20 +857,25 @@ function renderFunctionComponent( outlineModel(request, componentDebugInfo); emitDebugChunk(request, componentDebugID, componentDebugInfo); } - } - - prepareToUseHooksForComponent(prevThenableState, componentDebugInfo); - // The secondArg is always undefined in Server Components since refs error early. - const secondArg = undefined; - let result; - if (__DEV__) { + prepareToUseHooksForComponent(prevThenableState, componentDebugInfo); setCurrentOwner(componentDebugInfo); try { - result = Component(props, secondArg); + if (supportsComponentStorage) { + // Run the component in an Async Context that tracks the current owner. + result = componentStorage.run( + componentDebugInfo, + Component, + props, + secondArg, + ); + } else { + result = Component(props, secondArg); + } } finally { setCurrentOwner(null); } } else { + prepareToUseHooksForComponent(prevThenableState, null); result = Component(props, secondArg); } if (typeof result === 'object' && result !== null) { diff --git a/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js b/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js index 3ccd00b466459..f5f031a860ff9 100644 --- a/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js +++ b/packages/react-server/src/flight/ReactFlightAsyncDispatcher.js @@ -15,7 +15,7 @@ import {resolveRequest, getCache} from '../ReactFlightServer'; import {disableStringRefs} from 'shared/ReactFeatureFlags'; -import {currentOwner} from './ReactFlightCurrentOwner'; +import {resolveOwner} from './ReactFlightCurrentOwner'; function resolveCache(): Map { const request = resolveRequest(); @@ -39,9 +39,7 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({ }: any); if (__DEV__) { - DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => { - return currentOwner; - }; + DefaultAsyncDispatcher.getOwner = resolveOwner; } else if (!disableStringRefs) { // Server Components never use string refs but the JSX runtime looks for it. DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => { diff --git a/packages/react-server/src/flight/ReactFlightCurrentOwner.js b/packages/react-server/src/flight/ReactFlightCurrentOwner.js index 6b5d2dd46a8b4..fec9e86829ba4 100644 --- a/packages/react-server/src/flight/ReactFlightCurrentOwner.js +++ b/packages/react-server/src/flight/ReactFlightCurrentOwner.js @@ -9,8 +9,22 @@ import type {ReactComponentInfo} from 'shared/ReactTypes'; -export let currentOwner: ReactComponentInfo | null = null; +import { + supportsComponentStorage, + componentStorage, +} from '../ReactFlightServerConfig'; + +let currentOwner: ReactComponentInfo | null = null; export function setCurrentOwner(componentInfo: null | ReactComponentInfo) { currentOwner = componentInfo; } + +export function resolveOwner(): null | ReactComponentInfo { + if (currentOwner) return currentOwner; + if (supportsComponentStorage) { + const owner = componentStorage.getStore(); + if (owner) return owner; + } + return null; +}