[Flight] Add "use ..." boundary after the change instead of before it

fixtures/flight/src/App.js

@@ -33,20 +33,26 @@ function Foo({children}) {
   return <div>{children}</div>;
 }
 
+async function delay(text, ms) {
+  return new Promise(resolve => setTimeout(() => resolve(text), ms));
+}
+
 async function Bar({children}) {
-  await new Promise(resolve => setTimeout(() => resolve('deferred text'), 10));
+  await delay('deferred text', 10);
   return <div>{children}</div>;
 }
 
 async function ThirdPartyComponent() {
-  return new Promise(resolve =>
-    setTimeout(() => resolve('hello from a 3rd party'), 30)
-  );
+  return delay('hello from a 3rd party', 30);
 }
 
 // Using Web streams for tee'ing convenience here.
 let cachedThirdPartyReadableWeb;
 
+// We create the Component outside of AsyncLocalStorage so that it has no owner.
+// That way it gets the owner from the call to createFromNodeStream.
+const thirdPartyComponent = <ThirdPartyComponent />;
+
 function fetchThirdParty(noCache) {
   if (cachedThirdPartyReadableWeb && !noCache) {
     const [readableWeb1, readableWeb2] = cachedThirdPartyReadableWeb.tee();
@@ -59,7 +65,7 @@ function fetchThirdParty(noCache) {
   }
 
   const stream = renderToPipeableStream(
-    <ThirdPartyComponent />,
+    thirdPartyComponent,
     {},
     {environmentName: 'third-party'}
   );
@@ -80,8 +86,8 @@ function fetchThirdParty(noCache) {
 }
 
 async function ServerComponent({noCache}) {
-  await new Promise(resolve => setTimeout(() => resolve('deferred text'), 50));
-  return fetchThirdParty(noCache);
+  await delay('deferred text', 50);
+  return await fetchThirdParty(noCache);
 }
 
 export default async function App({prerender, noCache}) {

packages/react-client/src/ReactFlightClient.js

@@ -844,7 +844,7 @@ function createElement(
       // This owner should ideally have already been initialized to avoid getting
       // user stack frames on the stack.
       const ownerTask =
-        owner === null ? null : initializeFakeTask(response, owner, env);
+        owner === null ? null : initializeFakeTask(response, owner);
       if (ownerTask === null) {
         const rootTask = response._debugRootTask;
         if (rootTask != null) {
@@ -2494,7 +2494,6 @@ function getRootTask(
 function initializeFakeTask(
   response: Response,
   debugInfo: ReactComponentInfo | ReactAsyncInfo | ReactIOInfo,
-  childEnvironmentName: string,
 ): null | ConsoleTask {
   if (!supportsCreateTask) {
     return null;
@@ -2504,6 +2503,10 @@ function initializeFakeTask(
     // If it's null, we can't initialize a task.
     return null;
   }
+  const cachedEntry = debugInfo.debugTask;
+  if (cachedEntry !== undefined) {
+    return cachedEntry;
+  }
 
   // Workaround for a bug where Chrome Performance tracking uses the enclosing line/column
   // instead of the callsite. For ReactAsyncInfo/ReactIOInfo, the only thing we're going
@@ -2516,47 +2519,35 @@ function initializeFakeTask(
   const stack = debugInfo.stack;
   const env: string =
     debugInfo.env == null ? response._rootEnvironmentName : debugInfo.env;
-  if (env !== childEnvironmentName) {
+  const ownerEnv: string =
+    debugInfo.owner == null || debugInfo.owner.env == null
+      ? response._rootEnvironmentName
+      : debugInfo.owner.env;
+  const ownerTask =
+    debugInfo.owner == null
+      ? null
+      : initializeFakeTask(response, debugInfo.owner);
+  const taskName =
     // This is the boundary between two environments so we'll annotate the task name.
-    // That is unusual so we don't cache it.
-    const ownerTask =
-      debugInfo.owner == null
-        ? null
-        : initializeFakeTask(response, debugInfo.owner, env);
-    return buildFakeTask(
-      response,
-      ownerTask,
-      stack,
-      '"use ' + childEnvironmentName.toLowerCase() + '"',
-      env,
-      useEnclosingLine,
-    );
-  } else {
-    const cachedEntry = debugInfo.debugTask;
-    if (cachedEntry !== undefined) {
-      return cachedEntry;
-    }
-    const ownerTask =
-      debugInfo.owner == null
-        ? null
-        : initializeFakeTask(response, debugInfo.owner, env);
-    // Some unfortunate pattern matching to refine the type.
-    const taskName =
-      debugInfo.key !== undefined
+    // We assume that the stack frame of the entry into the new environment was done
+    // from the old environment. So we use the owner's environment as the current.
+    env !== ownerEnv
+      ? '"use ' + env.toLowerCase() + '"'
+      : // Some unfortunate pattern matching to refine the type.
+        debugInfo.key !== undefined
         ? getServerComponentTaskName(((debugInfo: any): ReactComponentInfo))
         : debugInfo.name !== undefined
           ? getIOInfoTaskName(((debugInfo: any): ReactIOInfo))
           : getAsyncInfoTaskName(((debugInfo: any): ReactAsyncInfo));
-    // $FlowFixMe[cannot-write]: We consider this part of initialization.
-    return (debugInfo.debugTask = buildFakeTask(
-      response,
-      ownerTask,
-      stack,
-      taskName,
-      env,
-      useEnclosingLine,
-    ));
-  }
+  // $FlowFixMe[cannot-write]: We consider this part of initialization.
+  return (debugInfo.debugTask = buildFakeTask(
+    response,
+    ownerTask,
+    stack,
+    taskName,
+    ownerEnv,
+    useEnclosingLine,
+  ));
 }
 
 function buildFakeTask(
@@ -2658,27 +2649,30 @@ function resolveDebugInfo(
       'resolveDebugInfo should never be called in production mode. This is a bug in React.',
     );
   }
-  // We eagerly initialize the fake task because this resolving happens outside any
-  // render phase so we're not inside a user space stack at this point. If we waited
-  // to initialize it when we need it, we might be inside user code.
-  const env =
-    debugInfo.env === undefined ? response._rootEnvironmentName : debugInfo.env;
   if (debugInfo.stack !== undefined) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe[incompatible-type]
       debugInfo;
-    initializeFakeTask(response, componentInfoOrAsyncInfo, env);
+    // We eagerly initialize the fake task because this resolving happens outside any
+    // render phase so we're not inside a user space stack at this point. If we waited
+    // to initialize it when we need it, we might be inside user code.
+    initializeFakeTask(response, componentInfoOrAsyncInfo);
   }
-  if (debugInfo.owner === null && response._debugRootOwner != null) {
+  if (debugInfo.owner == null && response._debugRootOwner != null) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe: By narrowing `owner` to `null`, we narrowed `debugInfo` to `ReactComponentInfo`
       debugInfo;
     // $FlowFixMe[cannot-write]
     componentInfoOrAsyncInfo.owner = response._debugRootOwner;
+    // We clear the parsed stack frames to indicate that it needs to be re-parsed from debugStack.
+    // $FlowFixMe[cannot-write]
+    componentInfoOrAsyncInfo.stack = null;
     // We override the stack if we override the owner since the stack where the root JSX
     // was created on the server isn't very useful but where the request was made is.
     // $FlowFixMe[cannot-write]
     componentInfoOrAsyncInfo.debugStack = response._debugRootStack;
+    // $FlowFixMe[cannot-write]
+    componentInfoOrAsyncInfo.debugTask = response._debugRootTask;
   } else if (debugInfo.stack !== undefined) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe[incompatible-type]
@@ -2738,7 +2732,7 @@ const replayConsoleWithCallStack = {
         bindToConsole(methodName, args, env),
       );
       if (owner != null) {
-        const task = initializeFakeTask(response, owner, env);
+        const task = initializeFakeTask(response, owner);
         initializeFakeStack(response, owner);
         if (task !== null) {
           task.run(callStack);
@@ -2812,10 +2806,8 @@ function resolveConsoleEntry(
 }
 
 function initializeIOInfo(response: Response, ioInfo: ReactIOInfo): void {
-  const env =
-    ioInfo.env === undefined ? response._rootEnvironmentName : ioInfo.env;
   if (ioInfo.stack !== undefined) {
-    initializeFakeTask(response, ioInfo, env);
+    initializeFakeTask(response, ioInfo);
     initializeFakeStack(response, ioInfo);
   }
   // Adjust the time to the current environment's time space.

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -320,7 +320,6 @@ describe('ReactFlight', () => {
                 name: 'Greeting',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   firstName: 'Seb',
@@ -364,7 +363,6 @@ describe('ReactFlight', () => {
                 name: 'Greeting',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   firstName: 'Seb',
@@ -2812,7 +2810,6 @@ describe('ReactFlight', () => {
                 name: 'ServerComponent',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   transport: expect.arrayContaining([]),
@@ -2834,7 +2831,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyComponent',
                 env: 'third-party',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -2851,7 +2847,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyLazyComponent',
                 env: 'third-party',
                 key: null,
-                owner: null,
                 stack: '    in myLazy (at **)\n    in lazyInitializer (at **)',
                 props: {},
               },
@@ -2867,7 +2862,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyFragmentComponent',
                 env: 'third-party',
                 key: '3',
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -2941,7 +2935,6 @@ describe('ReactFlight', () => {
                 name: 'ServerComponent',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   transport: expect.arrayContaining([]),
@@ -2961,7 +2954,6 @@ describe('ReactFlight', () => {
                 name: 'Keyed',
                 env: 'Server',
                 key: 'keyed',
-                owner: null,
                 stack: '    in ServerComponent (at **)',
                 props: {
                   children: {},
@@ -2980,7 +2972,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyAsyncIterableComponent',
                 env: 'third-party',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -3137,7 +3128,6 @@ describe('ReactFlight', () => {
                 name: 'Component',
                 env: 'A',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -3325,7 +3315,6 @@ describe('ReactFlight', () => {
           name: 'Greeting',
           env: 'Server',
           key: null,
-          owner: null,
           stack: '    in Object.<anonymous> (at **)',
           props: {
             firstName: 'Seb',

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js

@@ -708,7 +708,6 @@ describe('ReactFlightDOMBrowser', () => {
             name: 'Server',
             env: 'Server',
             key: null,
-            owner: null,
           }),
         }),
       );
@@ -724,7 +723,6 @@ describe('ReactFlightDOMBrowser', () => {
             name: 'Server',
             env: 'Server',
             key: null,
-            owner: null,
           }),
         }),
       );

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js

@@ -1190,7 +1190,6 @@ describe('ReactFlightDOMEdge', () => {
       const greetInfo = expect.objectContaining({
         name: 'Greeting',
         env: 'Server',
-        owner: null,
       });
       expect(lazyWrapper._debugInfo).toEqual([
         {time: 12},