experimental_use(promise) for Server Components (#25207)

Follow up to #25084. Implements experimental_use(promise) API in 
the Server Components runtime (Flight).

The implementation is much simpler than in Fiber because there is no
state. Even the "state" added in this PR — to track the result of each
promise across attempts — is reset as soon as a component 
successfully renders without suspending.

There are also fewer caveats around neglecting to cache a promise
because the state of the promises is preserved even if we switch to a
different task.

Server Components is the primary runtime where this API is intended to
be used.

The last runtime where we need to implement this is the server renderer
(Fizz).

Author: acdlite

packages/react-reconciler/src/ReactFiberWakeable.new.js

@@ -18,6 +18,7 @@ import type {
 let suspendedThenable: Thenable<mixed> | null = null;
 let adHocSuspendCount: number = 0;
 
+// TODO: Sparse arrays are bad for performance.
 let usedThenables: Array<Thenable<any> | void> | null = null;
 let lastUsedThenable: Thenable<any> | null = null;
 
@@ -74,6 +75,9 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
       suspendedThenable = null;
       break;
     default: {
+      // TODO: Only instrument the thenable if the status if not defined. If
+      // it's defined, but an unknown value, assume it's been instrumented by
+      // some custom userspace implementation.
       const pendingThenable: PendingThenable<mixed> = (thenable: any);
       pendingThenable.status = 'pending';
       pendingThenable.then(

packages/react-reconciler/src/ReactFiberWakeable.old.js

@@ -18,6 +18,7 @@ import type {
 let suspendedThenable: Thenable<mixed> | null = null;
 let adHocSuspendCount: number = 0;
 
+// TODO: Sparse arrays are bad for performance.
 let usedThenables: Array<Thenable<any> | void> | null = null;
 let lastUsedThenable: Thenable<any> | null = null;
 
@@ -74,6 +75,9 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
       suspendedThenable = null;
       break;
     default: {
+      // TODO: Only instrument the thenable if the status if not defined. If
+      // it's defined, but an unknown value, assume it's been instrumented by
+      // some custom userspace implementation.
       const pendingThenable: PendingThenable<mixed> = (thenable: any);
       pendingThenable.status = 'pending';
       pendingThenable.then(

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

@@ -24,6 +24,7 @@ let ReactDOMServer;
 let ReactServerDOMWriter;
 let ReactServerDOMReader;
 let Suspense;
+let use;
 
 describe('ReactFlightDOMBrowser', () => {
   beforeEach(() => {
@@ -39,6 +40,7 @@ describe('ReactFlightDOMBrowser', () => {
     ReactServerDOMWriter = require('react-server-dom-webpack/writer.browser.server');
     ReactServerDOMReader = require('react-server-dom-webpack');
     Suspense = React.Suspense;
+    use = React.experimental_use;
   });
 
   async function waitForSuspense(fn) {
@@ -562,4 +564,149 @@ describe('ReactFlightDOMBrowser', () => {
 
     expect(reportedErrors).toEqual(['for reasons']);
   });
+
+  // @gate enableUseHook
+  it('basic use(promise)', async () => {
+    function Server() {
+      return (
+        use(Promise.resolve('A')) +
+        use(Promise.resolve('B')) +
+        use(Promise.resolve('C'))
+      );
+    }
+
+    const stream = ReactServerDOMWriter.renderToReadableStream(<Server />);
+    const response = ReactServerDOMReader.createFromReadableStream(stream);
+
+    function Client() {
+      return response.readRoot();
+    }
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(
+        <Suspense fallback="Loading...">
+          <Client />
+        </Suspense>,
+      );
+    });
+    expect(container.innerHTML).toBe('ABC');
+  });
+
+  // @gate enableUseHook
+  it('use(promise) in multiple components', async () => {
+    function Child({prefix}) {
+      return prefix + use(Promise.resolve('C')) + use(Promise.resolve('D'));
+    }
+
+    function Parent() {
+      return (
+        <Child prefix={use(Promise.resolve('A')) + use(Promise.resolve('B'))} />
+      );
+    }
+
+    const stream = ReactServerDOMWriter.renderToReadableStream(<Parent />);
+    const response = ReactServerDOMReader.createFromReadableStream(stream);
+
+    function Client() {
+      return response.readRoot();
+    }
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(
+        <Suspense fallback="Loading...">
+          <Client />
+        </Suspense>,
+      );
+    });
+    expect(container.innerHTML).toBe('ABCD');
+  });
+
+  // @gate enableUseHook
+  it('using a rejected promise will throw', async () => {
+    const promiseA = Promise.resolve('A');
+    const promiseB = Promise.reject(new Error('Oops!'));
+    const promiseC = Promise.resolve('C');
+
+    // Jest/Node will raise an unhandled rejected error unless we await this. It
+    // works fine in the browser, though.
+    await expect(promiseB).rejects.toThrow('Oops!');
+
+    function Server() {
+      return use(promiseA) + use(promiseB) + use(promiseC);
+    }
+
+    const reportedErrors = [];
+    const stream = ReactServerDOMWriter.renderToReadableStream(
+      <Server />,
+      webpackMap,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
+    );
+    const response = ReactServerDOMReader.createFromReadableStream(stream);
+
+    class ErrorBoundary extends React.Component {
+      state = {error: null};
+      static getDerivedStateFromError(error) {
+        return {error};
+      }
+      render() {
+        if (this.state.error) {
+          return this.state.error.message;
+        }
+        return this.props.children;
+      }
+    }
+
+    function Client() {
+      return response.readRoot();
+    }
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(
+        <ErrorBoundary>
+          <Client />
+        </ErrorBoundary>,
+      );
+    });
+    expect(container.innerHTML).toBe('Oops!');
+    expect(reportedErrors.length).toBe(1);
+    expect(reportedErrors[0].message).toBe('Oops!');
+  });
+
+  // @gate enableUseHook
+  it("use a promise that's already been instrumented and resolved", async () => {
+    const thenable = {
+      status: 'fulfilled',
+      value: 'Hi',
+      then() {},
+    };
+
+    // This will never suspend because the thenable already resolved
+    function Server() {
+      return use(thenable);
+    }
+
+    const stream = ReactServerDOMWriter.renderToReadableStream(<Server />);
+    const response = ReactServerDOMReader.createFromReadableStream(stream);
+
+    function Client() {
+      return response.readRoot();
+    }
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<Client />);
+    });
+    expect(container.innerHTML).toBe('Hi');
+  });
 });

packages/react-server/src/ReactFlightHooks.js

@@ -9,11 +9,20 @@
 
 import type {Dispatcher as DispatcherType} from 'react-reconciler/src/ReactInternalTypes';
 import type {Request} from './ReactFlightServer';
-import type {ReactServerContext} from 'shared/ReactTypes';
+import type {ReactServerContext, Thenable, Usable} from 'shared/ReactTypes';
+import type {ThenableState} from './ReactFlightWakeable';
 import {REACT_SERVER_CONTEXT_TYPE} from 'shared/ReactSymbols';
 import {readContext as readContextImpl} from './ReactFlightNewContext';
+import {enableUseHook} from 'shared/ReactFeatureFlags';
+import {
+  getPreviouslyUsedThenableAtIndex,
+  createThenableState,
+  trackUsedThenable,
+} from './ReactFlightWakeable';
 
 let currentRequest = null;
+let thenableIndexCounter = 0;
+let thenableState = null;
 
 export function prepareToUseHooksForRequest(request: Request) {
   currentRequest = request;
@@ -23,6 +32,17 @@ export function resetHooksForRequest() {
   currentRequest = null;
 }
 
+export function prepareToUseHooksForComponent(
+  prevThenableState: ThenableState | null,
+) {
+  thenableIndexCounter = 0;
+  thenableState = prevThenableState;
+}
+
+export function getThenableStateAfterSuspending() {
+  return thenableState;
+}
+
 function readContext<T>(context: ReactServerContext<T>): T {
   if (__DEV__) {
     if (context.$$typeof !== REACT_SERVER_CONTEXT_TYPE) {
@@ -83,6 +103,7 @@ export const Dispatcher: DispatcherType = {
   useMemoCache(size: number): Array<any> {
     return new Array(size);
   },
+  use: enableUseHook ? use : (unsupportedHook: any),
 };
 
 function unsupportedHook(): void {
@@ -116,3 +137,72 @@ function useId(): string {
   // use 'S' for Flight components to distinguish from 'R' and 'r' in Fizz/Client
   return ':' + currentRequest.identifierPrefix + 'S' + id.toString(32) + ':';
 }
+
+function use<T>(usable: Usable<T>): T {
+  if (usable !== null && typeof usable === 'object') {
+    if (typeof usable.then === 'function') {
+      // This is a thenable.
+      const thenable: Thenable<T> = (usable: any);
+
+      // Track the position of the thenable within this fiber.
+      const index = thenableIndexCounter;
+      thenableIndexCounter += 1;
+
+      switch (thenable.status) {
+        case 'fulfilled': {
+          const fulfilledValue: T = thenable.value;
+          return fulfilledValue;
+        }
+        case 'rejected': {
+          const rejectedError = thenable.reason;
+          throw rejectedError;
+        }
+        default: {
+          const prevThenableAtIndex: Thenable<T> | null = getPreviouslyUsedThenableAtIndex(
+            thenableState,
+            index,
+          );
+          if (prevThenableAtIndex !== null) {
+            switch (prevThenableAtIndex.status) {
+              case 'fulfilled': {
+                const fulfilledValue: T = prevThenableAtIndex.value;
+                return fulfilledValue;
+              }
+              case 'rejected': {
+                const rejectedError: mixed = prevThenableAtIndex.reason;
+                throw rejectedError;
+              }
+              default: {
+                // The thenable still hasn't resolved. Suspend with the same
+                // thenable as last time to avoid redundant listeners.
+                throw prevThenableAtIndex;
+              }
+            }
+          } else {
+            // This is the first time something has been used at this index.
+            // Stash the thenable at the current index so we can reuse it during
+            // the next attempt.
+            if (thenableState === null) {
+              thenableState = createThenableState();
+            }
+            trackUsedThenable(thenableState, thenable, index);
+
+            // Suspend.
+            // TODO: Throwing here is an implementation detail that allows us to
+            // unwind the call stack. But we shouldn't allow it to leak into
+            // userspace. Throw an opaque placeholder value instead of the
+            // actual thenable. If it doesn't get captured by the work loop, log
+            // a warning, because that means something in userspace must have
+            // caught it.
+            throw thenable;
+          }
+        }
+      }
+    } else {
+      // TODO: Add support for Context
+    }
+  }
+
+  // eslint-disable-next-line react-internal/safe-string-coercion
+  throw new Error('An unsupported type was passed to use(): ' + String(usable));
+}