Fix lazy thenable suspension race

packages/react-reconciler/src/__tests__/ReactLazy-test.internal.js

@@ -123,18 +123,25 @@ describe('ReactLazy', () => {
       },
     }));
 
+    const LazyTextWithStatus = lazy(() => ({
+      status: 'fulfilled',
+      value: {default: Text},
+      then() {},
+    }));
+
     let root;
     await act(() => {
       root = ReactTestRenderer.create(
         <Suspense fallback={<Text text="Loading..." />}>
           <LazyText text="Hi" />
+          <LazyTextWithStatus text="Bye" />
         </Suspense>,
         {unstable_isConcurrent: true},
       );
     });
 
-    assertLog(['Hi']);
-    expect(root).toMatchRenderedOutput('Hi');
+    assertLog(['Hi', 'Bye']);
+    expect(root).toMatchRenderedOutput('HiBye');
   });
 
   it('can reject synchronously without suspending', async () => {
@@ -171,6 +178,76 @@ describe('ReactLazy', () => {
     expect(root).toMatchRenderedOutput('Error: oh no');
   });
 
+  it('can reject synchronously without suspending with status thenable', async () => {
+    const error = new Error('oh no');
+    const LazyText = lazy(() => ({
+      status: 'rejected',
+      reason: error,
+      then() {},
+    }));
+
+    class ErrorBoundary extends React.Component {
+      state = {};
+      static getDerivedStateFromError(error) {
+        return {message: error.message};
+      }
+      render() {
+        return this.state.message
+          ? `Error: ${this.state.message}`
+          : this.props.children;
+      }
+    }
+
+    let root;
+    await act(() => {
+      root = ReactTestRenderer.create(
+        <ErrorBoundary>
+          <Suspense fallback={<Text text="Loading..." />}>
+            <LazyText text="Hi" />
+          </Suspense>
+        </ErrorBoundary>,
+        {unstable_isConcurrent: true},
+      );
+    });
+    assertLog([]);
+    expect(root).toMatchRenderedOutput('Error: oh no');
+  });
+
+  it('can resolve after pending status thenable flips', async () => {
+    let statusThenable;
+    const LazyText = lazy(() => {
+      statusThenable = {
+        status: 'pending',
+        value: {default: Text},
+        then() {},
+      };
+      return statusThenable;
+    });
+
+    const root = ReactTestRenderer.create(
+      <Suspense fallback={<Text text="Loading..." />}>
+        <LazyText text="Hi" />
+      </Suspense>,
+      {
+        unstable_isConcurrent: true,
+      },
+    );
+
+    await waitForAll(['Loading...']);
+    expect(root).not.toMatchRenderedOutput('Hi');
+
+    statusThenable.status = 'fulfilled';
+    statusThenable.value = {default: Text};
+
+    root.update(
+      <Suspense fallback={<Text text="Loading..." />}>
+        <LazyText text="Hi again" />
+      </Suspense>,
+    );
+    await waitForAll(['Hi again']);
+    expect(root).toMatchRenderedOutput('Hi again');
+  });
+
   it('multiple lazy components', async () => {
     function Foo() {
       return <Text text="Foo" />;

packages/react/src/ReactLazy.js

@@ -31,24 +31,32 @@ type UninitializedPayload<T> = {
   _status: -1,
   _result: () => Thenable<{default: T, ...}>,
   _ioInfo?: ReactIOInfo, // DEV-only
+  _debugValueResolve?: null | ((value: mixed) => void),
+  _debugValueReject?: null | ((error: mixed) => void),
 };
 
 type PendingPayload = {
   _status: 0,
   _result: Wakeable,
   _ioInfo?: ReactIOInfo, // DEV-only
+  _debugValueResolve?: null | ((value: mixed) => void),
+  _debugValueReject?: null | ((error: mixed) => void),
 };
 
 type ResolvedPayload<T> = {
   _status: 1,
   _result: {default: T, ...},
   _ioInfo?: ReactIOInfo, // DEV-only
+  _debugValueResolve?: null | ((value: mixed) => void),
+  _debugValueReject?: null | ((error: mixed) => void),
 };
 
 type RejectedPayload = {
   _status: 2,
   _result: mixed,
   _ioInfo?: ReactIOInfo, // DEV-only
+  _debugValueResolve?: null | ((value: mixed) => void),
+  _debugValueReject?: null | ((error: mixed) => void),
 };
 
 type Payload<T> =
@@ -68,6 +76,62 @@ export type LazyComponent<T, P> = {
 };
 
 function lazyInitializer<T>(payload: Payload<T>): T {
+  if (payload._status === Pending) {
+    const thenable: Thenable<{default: T, ...}> = (payload._result: any);
+    switch (thenable.status) {
+      case 'fulfilled': {
+        const fulfilledThenable: FulfilledThenable<{default: T, ...}> =
+          (thenable: any);
+        const resolved: ResolvedPayload<T> = (payload: any);
+        resolved._status = Resolved;
+        resolved._result = fulfilledThenable.value;
+        if (__DEV__ && enableAsyncDebugInfo) {
+          const ioInfo = payload._ioInfo;
+          if (ioInfo != null) {
+            const ioInfoMutable: any = ioInfo;
+            ioInfoMutable.end = performance.now();
+            const debugValue =
+              fulfilledThenable.value == null
+                ? undefined
+                : fulfilledThenable.value.default;
+            const ioInfoValue: any = ioInfoMutable.value;
+            ioInfoValue.status = 'fulfilled';
+            ioInfoValue.value = debugValue;
+            const debugValueResolve = payload._debugValueResolve;
+            if (debugValueResolve != null) {
+              payload._debugValueResolve = null;
+              debugValueResolve(debugValue);
+            }
+          }
+        }
+        break;
+      }
+      case 'rejected': {
+        const rejectedThenable: RejectedThenable<{default: T, ...}> =
+          (thenable: any);
+        const rejected: RejectedPayload = (payload: any);
+        rejected._status = Rejected;
+        rejected._result = rejectedThenable.reason;
+        if (__DEV__ && enableAsyncDebugInfo) {
+          const ioInfo = payload._ioInfo;
+          if (ioInfo != null) {
+            const ioInfoMutable: any = ioInfo;
+            ioInfoMutable.end = performance.now();
+            const ioInfoValue: any = ioInfoMutable.value;
+            ioInfoValue.then(noop, noop);
+            ioInfoValue.status = 'rejected';
+            ioInfoValue.reason = rejectedThenable.reason;
+            const debugValueReject = payload._debugValueReject;
+            if (debugValueReject != null) {
+              payload._debugValueReject = null;
+              debugValueReject(rejectedThenable.reason);
+            }
+          }
+        }
+        break;
+      }
+    }
+  }
   if (payload._status === Uninitialized) {
     let resolveDebugValue: (void | T) => void = (null: any);
     let rejectDebugValue: mixed => void = (null: any);
@@ -83,6 +147,8 @@ function lazyInitializer<T>(payload: Payload<T>): T {
           resolveDebugValue = resolve;
           rejectDebugValue = reject;
         });
+        payload._debugValueResolve = resolveDebugValue;
+        payload._debugValueReject = rejectDebugValue;
       }
     }
     const ctor = payload._result;
@@ -176,6 +242,53 @@ function lazyInitializer<T>(payload: Payload<T>): T {
         }
       }
     }
+    if (payload._status === Uninitialized || payload._status === Pending) {
+      switch ((thenable: Thenable<{default: T, ...}>).status) {
+        case 'fulfilled': {
+          const fulfilledThenable: FulfilledThenable<{default: T, ...}> =
+            (thenable: any);
+          const resolved: ResolvedPayload<T> = (payload: any);
+          resolved._status = Resolved;
+          resolved._result = fulfilledThenable.value;
+          if (__DEV__ && enableAsyncDebugInfo) {
+            const ioInfo = payload._ioInfo;
+            if (ioInfo != null) {
+              const ioInfoMutable: any = ioInfo;
+              ioInfoMutable.end = performance.now();
+              const debugValue =
+                fulfilledThenable.value == null
+                  ? undefined
+                  : fulfilledThenable.value.default;
+              resolveDebugValue(debugValue);
+              const ioInfoValue: any = ioInfoMutable.value;
+              ioInfoValue.status = 'fulfilled';
+              ioInfoValue.value = debugValue;
+            }
+          }
+          break;
+        }
+        case 'rejected': {
+          const rejectedThenable: RejectedThenable<{default: T, ...}> =
+            (thenable: any);
+          const rejected: RejectedPayload = (payload: any);
+          rejected._status = Rejected;
+          rejected._result = rejectedThenable.reason;
+          if (__DEV__ && enableAsyncDebugInfo) {
+            const ioInfo = payload._ioInfo;
+            if (ioInfo != null) {
+              const ioInfoMutable: any = ioInfo;
+              ioInfoMutable.end = performance.now();
+              const ioInfoValue: any = ioInfoMutable.value;
+              ioInfoValue.then(noop, noop);
+              rejectDebugValue(rejectedThenable.reason);
+              ioInfoValue.status = 'rejected';
+              ioInfoValue.reason = rejectedThenable.reason;
+            }
+          }
+          break;
+        }
+      }
+    }
     if (payload._status === Uninitialized) {
       // In case, we're still uninitialized, then we're waiting for the thenable
       // to resolve. Set it as pending in the meantime.