Support AsyncIterables in Fizz

Author: sebmarkbage

packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js

@@ -3346,7 +3346,7 @@ describe('ReactDOMFizzServer', () => {
     ]);
   });
 
-  it('Supports iterable', async () => {
+  it('supports iterable', async () => {
     const Immutable = require('immutable');
 
     const mappedJSX = Immutable.fromJS([
@@ -3366,7 +3366,71 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
-  it('Supports bigint', async () => {
+  // @gate enableAsyncIterableChildren
+  it('supports async generator component', async () => {
+    async function* App() {
+      yield <span key="1">{await Promise.resolve('Hi')}</span>;
+      yield ' ';
+      yield <span key="2">{await Promise.resolve('World')}</span>;
+    }
+
+    await act(async () => {
+      const {pipe} = renderToPipeableStream(
+        <div>
+          <App />
+        </div>,
+      );
+      pipe(writable);
+    });
+
+    // Each act retries once which causes a new ping which schedules
+    // new work but only after the act has finished rendering.
+    await act(() => {});
+    await act(() => {});
+    await act(() => {});
+    await act(() => {});
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <span>Hi</span> <span>World</span>
+      </div>,
+    );
+  });
+
+  // @gate enableAsyncIterableChildren
+  it('supports async iterable children', async () => {
+    const iterable = {
+      async *[Symbol.asyncIterator]() {
+        yield <span key="1">{await Promise.resolve('Hi')}</span>;
+        yield ' ';
+        yield <span key="2">{await Promise.resolve('World')}</span>;
+      },
+    };
+
+    function App({children}) {
+      return <div>{children}</div>;
+    }
+
+    await act(() => {
+      const {pipe} = renderToPipeableStream(<App>{iterable}</App>);
+      pipe(writable);
+    });
+
+    // Each act retries once which causes a new ping which schedules
+    // new work but only after the act has finished rendering.
+    await act(() => {});
+    await act(() => {});
+    await act(() => {});
+    await act(() => {});
+
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <span>Hi</span> <span>World</span>
+      </div>,
+    );
+  });
+
+  it('supports bigint', async () => {
     await act(async () => {
       const {pipe} = ReactDOMFizzServer.renderToPipeableStream(
         <div>{10n}</div>,

packages/react-server/src/ReactFizzHooks.js

@@ -25,7 +25,11 @@ import type {TransitionStatus} from './ReactFizzConfig';
 
 import {readContext as readContextImpl} from './ReactFizzNewContext';
 import {getTreeId} from './ReactFizzTreeContext';
-import {createThenableState, trackUsedThenable} from './ReactFizzThenable';
+import {
+  createThenableState,
+  trackUsedThenable,
+  readPreviousThenable,
+} from './ReactFizzThenable';
 
 import {makeId, NotPendingTransition} from './ReactFizzConfig';
 import {createFastHash} from './ReactServerStreamConfig';
@@ -229,6 +233,13 @@ export function prepareToUseHooks(
   thenableState = prevThenableState;
 }
 
+export function prepareToUseThenableState(
+  prevThenableState: ThenableState | null,
+): void {
+  thenableIndexCounter = 0;
+  thenableState = prevThenableState;
+}
+
 export function finishHooks(
   Component: any,
   props: any,
@@ -765,6 +776,15 @@ export function unwrapThenable<T>(thenable: Thenable<T>): T {
   return trackUsedThenable(thenableState, thenable, index);
 }
 
+export function readPreviousThenableFromState<T>(): T | void {
+  const index = thenableIndexCounter;
+  thenableIndexCounter += 1;
+  if (thenableState === null) {
+    return undefined;
+  }
+  return readPreviousThenable(thenableState, index);
+}
+
 function unsupportedRefresh() {
   throw new Error('Cache cannot be refreshed during server rendering.');
 }

packages/react-server/src/ReactFizzServer.js

@@ -98,6 +98,7 @@ import {
 } from './ReactFizzNewContext';
 import {
   prepareToUseHooks,
+  prepareToUseThenableState,
   finishHooks,
   checkDidRenderIdHook,
   resetHooksState,
@@ -106,6 +107,7 @@ import {
   setCurrentResumableState,
   getThenableStateAfterSuspending,
   unwrapThenable,
+  readPreviousThenableFromState,
   getActionStateCount,
   getActionStateMatchingIndex,
 } from './ReactFizzHooks';
@@ -115,6 +117,7 @@ import {emptyTreeContext, pushTreeContext} from './ReactFizzTreeContext';
 
 import {
   getIteratorFn,
+  ASYNC_ITERATOR,
   REACT_ELEMENT_TYPE,
   REACT_PORTAL_TYPE,
   REACT_LAZY_TYPE,
@@ -144,6 +147,7 @@ import {
   enableRenderableContext,
   enableRefAsProp,
   disableDefaultPropsExceptForClasses,
+  enableAsyncIterableChildren,
 } from 'shared/ReactFeatureFlags';
 
 import assign from 'shared/assign';
@@ -2165,6 +2169,7 @@ function validateIterable(
       // as its direct child since we can recreate those by rerendering the component
       // as needed.
       const isGeneratorComponent =
+        childIndex === -1 && // Only the root child is valid
         task.componentStack !== null &&
         task.componentStack.tag === 1 && // FunctionComponent
         // $FlowFixMe[method-unbinding]
@@ -2197,6 +2202,43 @@ function validateIterable(
   }
 }
 
+function validateAsyncIterable(
+  task: Task,
+  iterable: AsyncIterable<any>,
+  childIndex: number,
+  iterator: AsyncIterator<any>,
+): void {
+  if (__DEV__) {
+    if (iterator === iterable) {
+      // We don't support rendering Generators as props because it's a mutation.
+      // See https://github.com/facebook/react/issues/12995
+      // We do support generators if they were created by a GeneratorFunction component
+      // as its direct child since we can recreate those by rerendering the component
+      // as needed.
+      const isGeneratorComponent =
+        childIndex === -1 && // Only the root child is valid
+        task.componentStack !== null &&
+        task.componentStack.tag === 1 && // FunctionComponent
+        // $FlowFixMe[method-unbinding]
+        Object.prototype.toString.call(task.componentStack.type) ===
+          '[object AsyncGeneratorFunction]' &&
+        // $FlowFixMe[method-unbinding]
+        Object.prototype.toString.call(iterator) === '[object AsyncGenerator]';
+      if (!isGeneratorComponent) {
+        if (!didWarnAboutGenerators) {
+          console.error(
+            'Using AsyncIterators as children is unsupported and will likely yield ' +
+              'unexpected results because enumerating a generator mutates it. ' +
+              'You can use an AsyncIterable that can iterate multiple times over ' +
+              'the same items.',
+          );
+        }
+        didWarnAboutGenerators = true;
+      }
+    }
+  }
+}
+
 function warnOnFunctionType(invalidChild: Function) {
   if (__DEV__) {
     const name = invalidChild.displayName || invalidChild.name || 'Component';
@@ -2327,20 +2369,83 @@ function renderNodeDestructive(
         // TODO: This is not great but I think it's inherent to the id
         // generation algorithm.
         let step = iterator.next();
-        // If there are not entries, we need to push an empty so we start by checking that.
         if (!step.done) {
           const children = [];
           do {
             children.push(step.value);
             step = iterator.next();
           } while (!step.done);
           renderChildrenArray(request, task, children, childIndex);
-          return;
         }
         return;
       }
     }
 
+    if (
+      enableAsyncIterableChildren &&
+      typeof (node: any)[ASYNC_ITERATOR] === 'function'
+    ) {
+      const iterator: AsyncIterator<ReactNodeList> = (node: any)[
+        ASYNC_ITERATOR
+      ]();
+      if (iterator) {
+        if (__DEV__) {
+          validateAsyncIterable(task, (node: any), childIndex, iterator);
+        }
+        // TODO: Update the task.node to be the iterator to avoid asking
+        // for new iterators, but we currently warn for rendering these
+        // so needs some refactoring to deal with the warning.
+
+        // We need to push a component stack because if this suspends, we'll pop a stack.
+        const previousComponentStack = task.componentStack;
+        task.componentStack = createBuiltInComponentStack(
+          task,
+          'AsyncIterable',
+        );
+
+        // Restore the thenable state before resuming.
+        const prevThenableState = task.thenableState;
+        task.thenableState = null;
+        prepareToUseThenableState(prevThenableState);
+
+        // We need to know how many total children are in this set, so that we
+        // can allocate enough id slots to acommodate them. So we must exhaust
+        // the iterator before we start recursively rendering the children.
+        // TODO: This is not great but I think it's inherent to the id
+        // generation algorithm.
+        const children = [];
+
+        let done = false;
+
+        if (iterator === node) {
+          // If it's an iterator we need to continue reading where we left
+          // off. We can do that by reading the first few rows from the previous
+          // thenable state.
+          // $FlowFixMe
+          let step = readPreviousThenableFromState();
+          while (step !== undefined) {
+            if (step.done) {
+              done = true;
+              break;
+            }
+            children.push(step.value);
+            step = readPreviousThenableFromState();
+          }
+        }
+
+        if (!done) {
+          let step = unwrapThenable(iterator.next());
+          while (!step.done) {
+            children.push(step.value);
+            step = unwrapThenable(iterator.next());
+          }
+        }
+        task.componentStack = previousComponentStack;
+        renderChildrenArray(request, task, children, childIndex);
+        return;
+      }
+    }
+
     // Usables are a valid React node type. When React encounters a Usable in
     // a child position, it unwraps it using the same algorithm as `use`. For
     // example, for promises, React will throw an exception to unwind the
@@ -3554,6 +3659,11 @@ function retryRenderTask(
         const ping = task.ping;
         x.then(ping, ping);
         task.thenableState = getThenableStateAfterSuspending();
+        // We pop one task off the stack because the node that suspended will be tried again,
+        // which will add it back onto the stack.
+        if (task.componentStack !== null) {
+          task.componentStack = task.componentStack.parent;
+        }
         return;
       } else if (
         enablePostpone &&
@@ -3639,6 +3749,11 @@ function retryReplayTask(request: Request, task: ReplayTask): void {
         const ping = task.ping;
         x.then(ping, ping);
         task.thenableState = getThenableStateAfterSuspending();
+        // We pop one task off the stack because the node that suspended will be tried again,
+        // which will add it back onto the stack.
+        if (task.componentStack !== null) {
+          task.componentStack = task.componentStack.parent;
+        }
         return;
       }
     }

packages/react-server/src/ReactFizzThenable.js

@@ -131,6 +131,19 @@ export function trackUsedThenable<T>(
   }
 }
 
+export function readPreviousThenable<T>(
+  thenableState: ThenableState,
+  index: number,
+): void | T {
+  const previous = thenableState[index];
+  if (previous === undefined) {
+    return undefined;
+  } else {
+    // We assume this has been resolved already.
+    return (previous: any).value;
+  }
+}
+
 // This is used to track the actual thenable that suspended so it can be
 // passed to the rest of the Suspense implementation — which, for historical
 // reasons, expects to receive a thenable.