@samitouri / QOS-React-1 / commits / d5c303427e

[Flight] Track Owner on AsyncLocalStorage When Available (#28807)

Stacked on #28798. Add another AsyncLocalStorage to the FlightServerConfig. This context tracks data on a per component level. Currently the only thing we track is the owner in DEV. AsyncLocalStorage around each component comes with a performance cost so we only do it DEV. It's not generally a particularly safe operation because you can't necessarily associate side-effects with a component based on execution scope. It can be a lazy initializer or cache():ed code etc. We also don't support string refs anymore for a reason. However, it's good enough for optional dev only information like the owner.

Sebastian Markbåge committed May 3, 2024 at 22:29 UTC d5c303427e077f8da8f231a621664654027440d6
15 files changed +189 -28
packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js
+71
@@ -21,6 +21,8 @@ if (typeof Blob === 'undefined') {
21 if (typeof File === 'undefined') {
22 global.File = require('buffer').File;
23 }
24 +// Patch for Edge environments for global scope
25 +global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
26
27 // Don't wait before processing work on the server.
28 // TODO: we can replace this with FlightServer.act().
@@ -32,6 +34,7 @@ let webpackMap;
34 let webpackModules;
35 let webpackModuleLoading;
36 let React;
37 +let ReactServer;
38 let ReactDOMServer;
39 let ReactServerDOMServer;
40 let ReactServerDOMClient;
@@ -55,6 +58,7 @@ describe('ReactFlightDOMEdge', () => {
58 webpackModules = WebpackMock.webpackModules;
59 webpackModuleLoading = WebpackMock.moduleLoading;
60
61 + ReactServer = require('react');
62 ReactServerDOMServer = require('react-server-dom-webpack/server');
63
64 jest.resetModules();
@@ -692,4 +696,71 @@ describe('ReactFlightDOMEdge', () => {
696 ),
697 );
698 });
699 +
700 + it('supports async server component debug info as the element owner in DEV', async () => {
701 + function Container({children}) {
702 + return children;
703 + }
704 +
705 + const promise = Promise.resolve(true);
706 + async function Greeting({firstName}) {
707 + // We can't use JSX here because it'll use the Client React.
708 + const child = ReactServer.createElement(
709 + 'span',
710 + null,
711 + 'Hello, ' + firstName,
712 + );
713 + // Yield the synchronous pass
714 + await promise;
715 + // We should still be able to track owner using AsyncLocalStorage.
716 + return ReactServer.createElement(Container, null, child);
717 + }
718 +
719 + const model = {
720 + greeting: ReactServer.createElement(Greeting, {firstName: 'Seb'}),
721 + };
722 +
723 + const stream = ReactServerDOMServer.renderToReadableStream(
724 + model,
725 + webpackMap,
726 + );
727 +
728 + const rootModel = await ReactServerDOMClient.createFromReadableStream(
729 + stream,
730 + {
731 + ssrManifest: {
732 + moduleMap: null,
733 + moduleLoading: null,
734 + },
735 + },
736 + );
737 +
738 + const ssrStream = await ReactDOMServer.renderToReadableStream(
739 + rootModel.greeting,
740 + );
741 + const result = await readResult(ssrStream);
742 + expect(result).toEqual('<span>Hello, Seb</span>');
743 +
744 + // Resolve the React Lazy wrapper which must have resolved by now.
745 + const lazyWrapper = rootModel.greeting;
746 + const greeting = lazyWrapper._init(lazyWrapper._payload);
747 +
748 + // We've rendered down to the span.
749 + expect(greeting.type).toBe('span');
750 + if (__DEV__) {
751 + const greetInfo = {name: 'Greeting', env: 'Server', owner: null};
752 + expect(lazyWrapper._debugInfo).toEqual([
753 + greetInfo,
754 + {name: 'Container', env: 'Server', owner: greetInfo},
755 + ]);
756 + // The owner that created the span was the outer server component.
757 + // We expect the debug info to be referentially equal to the owner.
758 + expect(greeting._owner).toBe(lazyWrapper._debugInfo[0]);
759 + } else {
760 + expect(lazyWrapper._debugInfo).toBe(undefined);
761 + expect(greeting._owner).toBe(
762 + gate(flags => flags.disableStringRefs) ? undefined : null,
763 + );
764 + }
765 + });
766 });
packages/react-server/src/ReactFlightServer.js
+24 -15
@@ -73,6 +73,8 @@ import {
73 isServerReference,
74 supportsRequestStorage,
75 requestStorage,
76 + supportsComponentStorage,
77 + componentStorage,
78 createHints,
79 initAsyncDebugInfo,
80 } from './ReactFlightServerConfig';
@@ -89,11 +91,9 @@ import {
91 getThenableStateAfterSuspending,
92 resetHooksForRequest,
93 } from './ReactFlightHooks';
92 -import {
93 - DefaultAsyncDispatcher,
94 - currentOwner,
95 - setCurrentOwner,
96 -} from './flight/ReactFlightAsyncDispatcher';
94 +import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
95 +
96 +import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
97
98 import {
99 getIteratorFn,
@@ -162,7 +162,7 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
162 // We don't currently use this id for anything but we emit it so that we can later
163 // refer to previous logs in debug info to associate them with a component.
164 const id = request.nextChunkId++;
165 - const owner: null | ReactComponentInfo = currentOwner;
165 + const owner: null | ReactComponentInfo = resolveOwner();
166 emitConsoleChunk(request, id, methodName, owner, stack, arguments);
167 }
168 // $FlowFixMe[prop-missing]
@@ -824,7 +824,11 @@ function renderFunctionComponent<Props>(
824 const prevThenableState = task.thenableState;
825 task.thenableState = null;
826
827 - let componentDebugInfo: null | ReactComponentInfo = null;
827 + // The secondArg is always undefined in Server Components since refs error early.
828 + const secondArg = undefined;
829 + let result;
830 +
831 + let componentDebugInfo: ReactComponentInfo;
832 if (__DEV__) {
833 if (debugID === null) {
834 // We don't have a chunk to assign debug info. We need to outline this
@@ -853,20 +857,25 @@ function renderFunctionComponent<Props>(
857 outlineModel(request, componentDebugInfo);
858 emitDebugChunk(request, componentDebugID, componentDebugInfo);
859 }
856 - }
857 -
858 - prepareToUseHooksForComponent(prevThenableState, componentDebugInfo);
859 - // The secondArg is always undefined in Server Components since refs error early.
860 - const secondArg = undefined;
861 - let result;
862 - if (__DEV__) {
860 + prepareToUseHooksForComponent(prevThenableState, componentDebugInfo);
861 setCurrentOwner(componentDebugInfo);
862 try {
865 - result = Component(props, secondArg);
863 + if (supportsComponentStorage) {
864 + // Run the component in an Async Context that tracks the current owner.
865 + result = componentStorage.run(
866 + componentDebugInfo,
867 + Component,
868 + props,
869 + secondArg,
870 + );
871 + } else {
872 + result = Component(props, secondArg);
873 + }
874 } finally {
875 setCurrentOwner(null);
876 }
877 } else {
878 + prepareToUseHooksForComponent(prevThenableState, null);
879 result = Component(props, secondArg);
880 }
881 if (typeof result === 'object' && result !== null) {
packages/react-server/src/flight/ReactFlightAsyncDispatcher.js
+3 -9
@@ -15,6 +15,8 @@ import {resolveRequest, getCache} from '../ReactFlightServer';
15
16 import {disableStringRefs} from 'shared/ReactFeatureFlags';
17
18 +import {resolveOwner} from './ReactFlightCurrentOwner';
19 +
20 function resolveCache(): Map<Function, mixed> {
21 const request = resolveRequest();
22 if (request) {
@@ -36,19 +38,11 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({
38 },
39 }: any);
40
39 -export let currentOwner: ReactComponentInfo | null = null;
40 -
41 if (__DEV__) {
42 - DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => {
43 - return currentOwner;
44 - };
42 + DefaultAsyncDispatcher.getOwner = resolveOwner;
43 } else if (!disableStringRefs) {
44 // Server Components never use string refs but the JSX runtime looks for it.
45 DefaultAsyncDispatcher.getOwner = (): null | ReactComponentInfo => {
46 return null;
47 };
48 }
51 -
52 -export function setCurrentOwner(componentInfo: null | ReactComponentInfo) {
53 - currentOwner = componentInfo;
54 -}
packages/react-server/src/flight/ReactFlightCurrentOwner.js new
+30
@@ -0,0 +1,30 @@
1 +/**
2 + * Copyright (c) Meta Platforms, Inc. and affiliates.
3 + *
4 + * This source code is licensed under the MIT license found in the
5 + * LICENSE file in the root directory of this source tree.
6 + *
7 + * @flow
8 + */
9 +
10 +import type {ReactComponentInfo} from 'shared/ReactTypes';
11 +
12 +import {
13 + supportsComponentStorage,
14 + componentStorage,
15 +} from '../ReactFlightServerConfig';
16 +
17 +let currentOwner: ReactComponentInfo | null = null;
18 +
19 +export function setCurrentOwner(componentInfo: null | ReactComponentInfo) {
20 + currentOwner = componentInfo;
21 +}
22 +
23 +export function resolveOwner(): null | ReactComponentInfo {
24 + if (currentOwner) return currentOwner;
25 + if (supportsComponentStorage) {
26 + const owner = componentStorage.getStore();
27 + if (owner) return owner;
28 + }
29 + return null;
30 +}
packages/react-server/src/forks/ReactFlightServerConfig.custom.js
+5
@@ -8,6 +8,7 @@
8 */
9
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from '../ReactFlightServerConfigBundlerCustom';
14
@@ -23,6 +24,10 @@ export const isPrimaryRenderer = false;
24 export const supportsRequestStorage = false;
25 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
26
27 +export const supportsComponentStorage = false;
28 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
29 + (null: any);
30 +
31 export function createHints(): any {
32 return null;
33 }
packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-esm.js
+7 -4
@@ -6,15 +6,18 @@
6 *
7 * @flow
8 */
9 -import {AsyncLocalStorage} from 'async_hooks';
9
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from 'react-server-dom-esm/src/ReactFlightServerConfigESMBundler';
14 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
15
16 -export const supportsRequestStorage = true;
17 -export const requestStorage: AsyncLocalStorage<Request | void> =
18 - new AsyncLocalStorage();
16 +export const supportsRequestStorage = false;
17 +export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
18 +
19 +export const supportsComponentStorage = false;
20 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
21 + (null: any);
22
23 export * from '../ReactFlightServerConfigDebugNoop';
packages/react-server/src/forks/ReactFlightServerConfig.dom-browser-turbopack.js
+5
@@ -8,6 +8,7 @@
8 */
9
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler';
14 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
16 export const supportsRequestStorage = false;
17 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
18
19 +export const supportsComponentStorage = false;
20 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
21 + (null: any);
22 +
23 export * from '../ReactFlightServerConfigDebugNoop';
packages/react-server/src/forks/ReactFlightServerConfig.dom-browser.js
+5
@@ -8,6 +8,7 @@
8 */
9
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler';
14 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
16 export const supportsRequestStorage = false;
17 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
18
19 +export const supportsComponentStorage = false;
20 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
21 + (null: any);
22 +
23 export * from '../ReactFlightServerConfigDebugNoop';
packages/react-server/src/forks/ReactFlightServerConfig.dom-bun.js
+5
@@ -8,6 +8,7 @@
8 */
9
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from '../ReactFlightServerConfigBundlerCustom';
14 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
16 export const supportsRequestStorage = false;
17 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
18
19 +export const supportsComponentStorage = false;
20 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
21 + (null: any);
22 +
23 export * from '../ReactFlightServerConfigDebugNoop';
packages/react-server/src/forks/ReactFlightServerConfig.dom-edge-turbopack.js
+6
@@ -7,6 +7,7 @@
7 * @flow
8 */
9 import type {Request} from 'react-server/src/ReactFlightServer';
10 +import type {ReactComponentInfo} from 'shared/ReactTypes';
11
12 export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler';
13 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -16,6 +17,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
17 export const requestStorage: AsyncLocalStorage<Request | void> =
18 supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
19
20 +export const supportsComponentStorage: boolean =
21 + __DEV__ && supportsRequestStorage;
22 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
23 + supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
24 +
25 // We use the Node version but get access to async_hooks from a global.
26 import type {HookCallbacks, AsyncHook} from 'async_hooks';
27 export const createAsyncHook: HookCallbacks => AsyncHook =
packages/react-server/src/forks/ReactFlightServerConfig.dom-edge.js
+7
@@ -6,7 +6,9 @@
6 *
7 * @flow
8 */
9 +
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler';
14 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -16,6 +18,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
18 export const requestStorage: AsyncLocalStorage<Request | void> =
19 supportsRequestStorage ? new AsyncLocalStorage() : (null: any);
20
21 +export const supportsComponentStorage: boolean =
22 + __DEV__ && supportsRequestStorage;
23 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
24 + supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
25 +
26 // We use the Node version but get access to async_hooks from a global.
27 import type {HookCallbacks, AsyncHook} from 'async_hooks';
28 export const createAsyncHook: HookCallbacks => AsyncHook =
packages/react-server/src/forks/ReactFlightServerConfig.dom-legacy.js
+5
@@ -8,6 +8,7 @@
8 */
9
10 import type {Request} from 'react-server/src/ReactFlightServer';
11 +import type {ReactComponentInfo} from 'shared/ReactTypes';
12
13 export * from '../ReactFlightServerConfigBundlerCustom';
14 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -15,4 +16,8 @@ export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
16 export const supportsRequestStorage = false;
17 export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);
18
19 +export const supportsComponentStorage = false;
20 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
21 + (null: any);
22 +
23 export * from '../ReactFlightServerConfigDebugNoop';
packages/react-server/src/forks/ReactFlightServerConfig.dom-node-esm.js
+6
@@ -6,9 +6,11 @@
6 *
7 * @flow
8 */
9 +
10 import {AsyncLocalStorage} from 'async_hooks';
11
12 import type {Request} from 'react-server/src/ReactFlightServer';
13 +import type {ReactComponentInfo} from 'shared/ReactTypes';
14
15 export * from 'react-server-dom-esm/src/ReactFlightServerConfigESMBundler';
16 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -17,5 +19,9 @@ export const supportsRequestStorage = true;
19 export const requestStorage: AsyncLocalStorage<Request | void> =
20 new AsyncLocalStorage();
21
22 +export const supportsComponentStorage = __DEV__;
23 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
24 + supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
25 +
26 export {createHook as createAsyncHook, executionAsyncId} from 'async_hooks';
27 export * from '../ReactFlightServerConfigDebugNode';
packages/react-server/src/forks/ReactFlightServerConfig.dom-node-turbopack.js
+5
@@ -10,6 +10,7 @@
10 import {AsyncLocalStorage} from 'async_hooks';
11
12 import type {Request} from 'react-server/src/ReactFlightServer';
13 +import type {ReactComponentInfo} from 'shared/ReactTypes';
14
15 export * from 'react-server-dom-turbopack/src/ReactFlightServerConfigTurbopackBundler';
16 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -18,5 +19,9 @@ export const supportsRequestStorage = true;
19 export const requestStorage: AsyncLocalStorage<Request | void> =
20 new AsyncLocalStorage();
21
22 +export const supportsComponentStorage = __DEV__;
23 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
24 + supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
25 +
26 export {createHook as createAsyncHook, executionAsyncId} from 'async_hooks';
27 export * from '../ReactFlightServerConfigDebugNode';
packages/react-server/src/forks/ReactFlightServerConfig.dom-node.js
+5
@@ -10,6 +10,7 @@
10 import {AsyncLocalStorage} from 'async_hooks';
11
12 import type {Request} from 'react-server/src/ReactFlightServer';
13 +import type {ReactComponentInfo} from 'shared/ReactTypes';
14
15 export * from 'react-server-dom-webpack/src/ReactFlightServerConfigWebpackBundler';
16 export * from 'react-dom-bindings/src/server/ReactFlightServerConfigDOM';
@@ -18,5 +19,9 @@ export const supportsRequestStorage = true;
19 export const requestStorage: AsyncLocalStorage<Request | void> =
20 new AsyncLocalStorage();
21
22 +export const supportsComponentStorage = __DEV__;
23 +export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
24 + supportsComponentStorage ? new AsyncLocalStorage() : (null: any);
25 +
26 export {createHook as createAsyncHook, executionAsyncId} from 'async_hooks';
27 export * from '../ReactFlightServerConfigDebugNode';