From d38ec17b1d6ab75715c4ec06f2793c2986dcd549 Mon Sep 17 00:00:00 2001
From: Joseph Savona <josephsavona@users.noreply.github.com>
Date: Mon, 7 Sep 2020 19:41:22 -0400
Subject: [PATCH] [Flight] Set dispatcher for duration of performWork()
 (#19776)

---
 packages/react-server/src/ReactFlightServer.js |  9 +++++----
 .../ReactFlightDOMRelay-test.internal.js       | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index 17058a9ac30ca..17fd077c57f88 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/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;
diff --git a/packages/react-transport-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js b/packages/react-transport-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js
index 67db8e74da9dd..756f1f5467ed6 100644
--- a/packages/react-transport-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js
+++ b/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);
+  });
 });