From e99bfd4c6d2222ea110aee694338ff1dc58e72c8 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Mon, 29 Mar 2021 15:39:05 -0400
Subject: [PATCH] Add fatal error handling

---
 .../ReactFlightDOMRelayServerHostConfig.js    |  7 +++-
 .../ReactFlightNativeRelayServerHostConfig.js |  7 +++-
 .../react-server/src/ReactFlightServer.js     | 40 +++++++++++++------
 .../src/ReactFlightServerConfigStream.js      |  1 +
 4 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js b/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js
index 1aa02e00d54b5..7ebcaf6b1e8d5 100644
--- a/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js
+++ b/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js
@@ -26,6 +26,7 @@ import {resolveModelToJSON} from 'react-server/src/ReactFlightServer';
 import {
   emitRow,
   resolveModuleMetaData as resolveModuleMetaDataImpl,
+  close,
 } from 'ReactFlightDOMRelayServerIntegration';
 
 export type {
@@ -146,4 +147,8 @@ export function writeChunk(destination: Destination, chunk: Chunk): boolean {
 
 export function completeWriting(destination: Destination) {}
 
-export {close} from 'ReactFlightDOMRelayServerIntegration';
+export {close};
+
+export function closeWithError(destination: Destination, error: mixed): void {
+  close(destination);
+}
diff --git a/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js b/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js
index ecea013654f39..8bf0cdf8b41ef 100644
--- a/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js
+++ b/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js
@@ -25,6 +25,7 @@ import {resolveModelToJSON} from 'react-server/src/ReactFlightServer';
 
 import {
   emitRow,
+  close,
   resolveModuleMetaData as resolveModuleMetaDataImpl,
 } from 'ReactFlightNativeRelayServerIntegration';
 
@@ -146,4 +147,8 @@ export function writeChunk(destination: Destination, chunk: Chunk): boolean {
 
 export function completeWriting(destination: Destination) {}
 
-export {close} from 'ReactFlightNativeRelayServerIntegration';
+export {close};
+
+export function closeWithError(destination: Destination, error: mixed): void {
+  close(destination);
+}
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index bb3c0307be877..2a09e21c89a00 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -24,6 +24,7 @@ import {
   completeWriting,
   flushBuffered,
   close,
+  closeWithError,
   processModelChunk,
   processModuleChunk,
   processSymbolChunk,
@@ -599,6 +600,11 @@ function reportError(request: Request, error: mixed): void {
   request.onError(error);
 }
 
+function fatalError(request: Request, error: mixed): void {
+  // This is called outside error handling code such as if an error happens in React internals.
+  closeWithError(request.destination, error);
+}
+
 function emitErrorChunk(request: Request, id: number, error: mixed): void {
   // TODO: We should not leak error messages to the client in prod.
   // Give this an error code instead and log on the server.
@@ -677,18 +683,23 @@ function performWork(request: Request): void {
   ReactCurrentDispatcher.current = Dispatcher;
   currentCache = request.cache;
 
-  const pingedSegments = request.pingedSegments;
-  request.pingedSegments = [];
-  for (let i = 0; i < pingedSegments.length; i++) {
-    const segment = pingedSegments[i];
-    retrySegment(request, segment);
-  }
-  if (request.flowing) {
-    flushCompletedChunks(request);
+  try {
+    const pingedSegments = request.pingedSegments;
+    request.pingedSegments = [];
+    for (let i = 0; i < pingedSegments.length; i++) {
+      const segment = pingedSegments[i];
+      retrySegment(request, segment);
+    }
+    if (request.flowing) {
+      flushCompletedChunks(request);
+    }
+  } catch (error) {
+    reportError(request, error);
+    fatalError(request, error);
+  } finally {
+    ReactCurrentDispatcher.current = prevDispatcher;
+    currentCache = prevCache;
   }
-
-  ReactCurrentDispatcher.current = prevDispatcher;
-  currentCache = prevCache;
 }
 
 let reentrant = false;
@@ -760,7 +771,12 @@ export function startWork(request: Request): void {
 
 export function startFlowing(request: Request): void {
   request.flowing = true;
-  flushCompletedChunks(request);
+  try {
+    flushCompletedChunks(request);
+  } catch (error) {
+    reportError(request, error);
+    fatalError(request, error);
+  }
 }
 
 function unsupportedHook(): void {
diff --git a/packages/react-server/src/ReactFlightServerConfigStream.js b/packages/react-server/src/ReactFlightServerConfigStream.js
index b9a57c16eb277..4ff841bfc11ba 100644
--- a/packages/react-server/src/ReactFlightServerConfigStream.js
+++ b/packages/react-server/src/ReactFlightServerConfigStream.js
@@ -126,4 +126,5 @@ export {
   writeChunk,
   completeWriting,
   close,
+  closeWithError,
 } from './ReactServerStreamConfig';