Skip to content

Commit

Permalink
[Flight] Track Owner on AsyncLocalStorage When Available (#28807)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
sebmarkbage committed May 3, 2024
1 parent 0a0a3af commit d5c3034
Show file tree
Hide file tree
Showing 15 changed files with 189 additions and 28 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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().
Expand All @@ -32,6 +34,7 @@ let webpackMap;
let webpackModules;
let webpackModuleLoading;
let React;
let ReactServer;
let ReactDOMServer;
let ReactServerDOMServer;
let ReactServerDOMClient;
Expand All @@ -55,6 +58,7 @@ describe('ReactFlightDOMEdge', () => {
webpackModules = WebpackMock.webpackModules;
webpackModuleLoading = WebpackMock.moduleLoading;

ReactServer = require('react');
ReactServerDOMServer = require('react-server-dom-webpack/server');

jest.resetModules();
Expand Down Expand Up @@ -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('<span>Hello, Seb</span>');

// 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,
);
}
});
});
39 changes: 24 additions & 15 deletions packages/react-server/src/ReactFlightServer.js
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ import {
isServerReference,
supportsRequestStorage,
requestStorage,
supportsComponentStorage,
componentStorage,
createHints,
initAsyncDebugInfo,
} from './ReactFlightServerConfig';
Expand All @@ -89,11 +91,9 @@ import {
getThenableStateAfterSuspending,
resetHooksForRequest,
} from './ReactFlightHooks';
import {
DefaultAsyncDispatcher,
currentOwner,
setCurrentOwner,
} from './flight/ReactFlightAsyncDispatcher';
import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';

import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';

import {
getIteratorFn,
Expand Down Expand Up @@ -162,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]
Expand Down Expand Up @@ -824,7 +824,11 @@ function renderFunctionComponent<Props>(
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
Expand Down Expand Up @@ -853,20 +857,25 @@ function renderFunctionComponent<Props>(
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) {
Expand Down
12 changes: 3 additions & 9 deletions packages/react-server/src/flight/ReactFlightAsyncDispatcher.js
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ import {resolveRequest, getCache} from '../ReactFlightServer';

import {disableStringRefs} from 'shared/ReactFeatureFlags';

import {resolveOwner} from './ReactFlightCurrentOwner';

function resolveCache(): Map<Function, mixed> {
const request = resolveRequest();
if (request) {
Expand All @@ -36,19 +38,11 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({
},
}: any);

export let currentOwner: ReactComponentInfo | null = null;

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 => {
return null;
};
}

export function setCurrentOwner(componentInfo: null | ReactComponentInfo) {
currentOwner = componentInfo;
}
30 changes: 30 additions & 0 deletions packages/react-server/src/flight/ReactFlightCurrentOwner.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
/**
* 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';

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;
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
*/

import type {Request} from 'react-server/src/ReactFlightServer';
import type {ReactComponentInfo} from 'shared/ReactTypes';

export * from '../ReactFlightServerConfigBundlerCustom';

Expand All @@ -23,6 +24,10 @@ export const isPrimaryRenderer = false;
export const supportsRequestStorage = false;
export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

export const supportsComponentStorage = false;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
(null: any);

export function createHints(): any {
return null;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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<Request | void> =
new AsyncLocalStorage();
export const supportsRequestStorage = false;
export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

export const supportsComponentStorage = false;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
(null: any);

export * from '../ReactFlightServerConfigDebugNoop';
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,16 @@
*/

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';

export const supportsRequestStorage = false;
export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

export const supportsComponentStorage = false;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
(null: any);

export * from '../ReactFlightServerConfigDebugNoop';
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,16 @@
*/

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';

export const supportsRequestStorage = false;
export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

export const supportsComponentStorage = false;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
(null: any);

export * from '../ReactFlightServerConfigDebugNoop';
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,16 @@
*/

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';

export const supportsRequestStorage = false;
export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

export const supportsComponentStorage = false;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
(null: any);

export * from '../ReactFlightServerConfigDebugNoop';
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -16,6 +17,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
export const requestStorage: AsyncLocalStorage<Request | void> =
supportsRequestStorage ? new AsyncLocalStorage() : (null: any);

export const supportsComponentStorage: boolean =
__DEV__ && supportsRequestStorage;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
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 =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -16,6 +18,11 @@ export const supportsRequestStorage = typeof AsyncLocalStorage === 'function';
export const requestStorage: AsyncLocalStorage<Request | void> =
supportsRequestStorage ? new AsyncLocalStorage() : (null: any);

export const supportsComponentStorage: boolean =
__DEV__ && supportsRequestStorage;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
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 =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,16 @@
*/

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';

export const supportsRequestStorage = false;
export const requestStorage: AsyncLocalStorage<Request | void> = (null: any);

export const supportsComponentStorage = false;
export const componentStorage: AsyncLocalStorage<ReactComponentInfo | void> =
(null: any);

export * from '../ReactFlightServerConfigDebugNoop';

0 comments on commit d5c3034

Please sign in to comment.