[Flight] Set dispatcher for duration of performWork()

packages/react-server/src/ReactFlightServer.js

@@ -273,11 +273,9 @@ export function resolveModelToJSON(
     value !== null &&
     value.$$typeof === REACT_ELEMENT_TYPE
   ) {
-    const prevDispatcher = ReactCurrentDispatcher.current;
     // TODO: Concatenate keys of parents onto children.
     const element: React$Element<any> = (value: any);
     try {
-      ReactCurrentDispatcher.current = Dispatcher;
       // Attempt to render the server component.
       value = attemptResolveElement(element);
     } catch (x) {
@@ -292,8 +290,6 @@ export function resolveModelToJSON(
         // Something errored. Don't bother encoding anything up to here.
         throw x;
       }
-    } finally {
-      ReactCurrentDispatcher.current = prevDispatcher;
     }
   }
 
@@ -355,6 +351,9 @@ function retrySegment(request: Request, segment: Segment): void {
 }
 
 function performWork(request: Request): void {
+  const prevDispatcher = ReactCurrentDispatcher.current;
+  ReactCurrentDispatcher.current = Dispatcher;
+
   const pingedSegments = request.pingedSegments;
   request.pingedSegments = [];
   for (let i = 0; i < pingedSegments.length; i++) {
@@ -364,6 +363,8 @@ function performWork(request: Request): void {
   if (request.flowing) {
     flushCompletedChunks(request);
   }
+
+  ReactCurrentDispatcher.current = prevDispatcher;
 }
 
 let reentrant = false;

packages/react-transport-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js

@@ -213,4 +213,22 @@ describe('ReactFlightDOMRelay', () => {
       },
     });
   });
+
+  it('can handle a subset of Hooks, with element as root', () => {
+    const {useMemo, useCallback} = React;
+    function Inner({x}) {
+      const foo = useMemo(() => x + x, [x]);
+      const bar = useCallback(() => 10 + foo, [foo]);
+      return bar();
+    }
+
+    function Foo() {
+      return <Inner x={2} />;
+    }
+    const transport = [];
+    ReactDOMFlightRelayServer.render(<Foo />, transport);
+
+    const model = readThrough(transport);
+    expect(model).toEqual(14);
+  });
 });