Abort Flight (#24754)

sebmarkbage · web-flow · commit 56389e81ffe1 · 2022-06-19T11:05:41.000-04:00
Add aborting to the Flight Server. This encodes the reason as an "error"
row that gets thrown client side. These are still exposed in prod which
is a follow up we'll still have to do to encode them as digests instead.

The error is encoded once and then referenced by each row that needs to
be updated.
diff --git a/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js b/packages/react-server-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js
@@ -117,6 +117,14 @@ export function processModelChunk(
   return ['J', id, json];
 }
 
+export function processReferenceChunk(
+  request: Request,
+  id: number,
+  reference: string,
+): Chunk {
+  return ['J', id, reference];
+}
+
 export function processModuleChunk(
   request: Request,
   id: number,
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js
@@ -15,12 +15,14 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  abort,
 } from 'react-server/src/ReactFlightServer';
 
 type Options = {
-  onError?: (error: mixed) => void,
-  context?: Array<[string, ServerContextJSONValue]>,
   identifierPrefix?: string,
+  signal?: AbortSignal,
+  context?: Array<[string, ServerContextJSONValue]>,
+  onError?: (error: mixed) => void,
 };
 
 function renderToReadableStream(
@@ -35,6 +37,18 @@ function renderToReadableStream(
     options ? options.context : undefined,
     options ? options.identifierPrefix : undefined,
   );
+  if (options && options.signal) {
+    const signal = options.signal;
+    if (signal.aborted) {
+      abort(request, (signal: any).reason);
+    } else {
+      const listener = () => {
+        abort(request, (signal: any).reason);
+        signal.removeEventListener('abort', listener);
+      };
+      signal.addEventListener('abort', listener);
+    }
+  }
   const stream = new ReadableStream(
     {
       type: 'bytes',
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js
@@ -16,6 +16,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  abort,
 } from 'react-server/src/ReactFlightServer';
 
 function createDrainHandler(destination, request) {
@@ -29,6 +30,7 @@ type Options = {
 };
 
 type PipeableStream = {|
+  abort(reason: mixed): void,
   pipe<T: Writable>(destination: T): T,
 |};
 
@@ -58,6 +60,9 @@ function renderToPipeableStream(
       destination.on('drain', createDrainHandler(destination, request));
       return destination;
     },
+    abort(reason: mixed) {
+      abort(request, reason);
+    },
   };
 }
 
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
@@ -30,6 +30,7 @@ let React;
 let ReactDOMClient;
 let ReactServerDOMWriter;
 let ReactServerDOMReader;
+let Suspense;
 
 describe('ReactFlightDOM', () => {
   beforeEach(() => {
@@ -42,6 +43,7 @@ describe('ReactFlightDOM', () => {
     ReactDOMClient = require('react-dom/client');
     ReactServerDOMWriter = require('react-server-dom-webpack/writer.node.server');
     ReactServerDOMReader = require('react-server-dom-webpack');
+    Suspense = React.Suspense;
   });
 
   function getTestStream() {
@@ -92,6 +94,11 @@ describe('ReactFlightDOM', () => {
     }
   }
 
+  const theInfinitePromise = new Promise(() => {});
+  function InfiniteSuspend() {
+    throw theInfinitePromise;
+  }
+
   it('should resolve HTML using Node streams', async () => {
     function Text({children}) {
       return <span>{children}</span>;
@@ -133,8 +140,6 @@ describe('ReactFlightDOM', () => {
   });
 
   it('should resolve the root', async () => {
-    const {Suspense} = React;
-
     // Model
     function Text({children}) {
       return <span>{children}</span>;
@@ -184,8 +189,6 @@ describe('ReactFlightDOM', () => {
   });
 
   it('should not get confused by $', async () => {
-    const {Suspense} = React;
-
     // Model
     function RootModel() {
       return {text: '$1'};
@@ -220,8 +223,6 @@ describe('ReactFlightDOM', () => {
   });
 
   it('should not get confused by @', async () => {
-    const {Suspense} = React;
-
     // Model
     function RootModel() {
       return {text: '@div'};
@@ -257,7 +258,6 @@ describe('ReactFlightDOM', () => {
 
   it('should progressively reveal server components', async () => {
     let reportedErrors = [];
-    const {Suspense} = React;
 
     // Client Components
 
@@ -460,8 +460,6 @@ describe('ReactFlightDOM', () => {
   });
 
   it('should preserve state of client components on refetch', async () => {
-    const {Suspense} = React;
-
     // Client
 
     function Page({response}) {
@@ -545,4 +543,64 @@ describe('ReactFlightDOM', () => {
     expect(inputB.tagName).toBe('INPUT');
     expect(inputB.value).toBe('goodbye');
   });
+
+  it('should be able to complete after aborting and throw the reason client-side', async () => {
+    const reportedErrors = [];
+
+    class ErrorBoundary extends React.Component {
+      state = {hasError: false, error: null};
+      static getDerivedStateFromError(error) {
+        return {
+          hasError: true,
+          error,
+        };
+      }
+      render() {
+        if (this.state.hasError) {
+          return this.props.fallback(this.state.error);
+        }
+        return this.props.children;
+      }
+    }
+
+    const {writable, readable} = getTestStream();
+    const {pipe, abort} = ReactServerDOMWriter.renderToPipeableStream(
+      <div>
+        <InfiniteSuspend />
+      </div>,
+      webpackMap,
+      {
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
+    );
+    pipe(writable);
+    const response = ReactServerDOMReader.createFromReadableStream(readable);
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+
+    function App({res}) {
+      return res.readRoot();
+    }
+
+    await act(async () => {
+      root.render(
+        <ErrorBoundary fallback={e => <p>{e.message}</p>}>
+          <Suspense fallback={<p>(loading)</p>}>
+            <App res={response} />
+          </Suspense>
+        </ErrorBoundary>,
+      );
+    });
+    expect(container.innerHTML).toBe('<p>(loading)</p>');
+
+    await act(async () => {
+      abort('for reasons');
+    });
+    expect(container.innerHTML).toBe('<p>Error: for reasons</p>');
+
+    expect(reportedErrors).toEqual(['for reasons']);
+  });
 });
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js
@@ -27,6 +27,7 @@ let ReactDOMClient;
 let ReactDOMServer;
 let ReactServerDOMWriter;
 let ReactServerDOMReader;
+let Suspense;
 
 describe('ReactFlightDOMBrowser', () => {
   beforeEach(() => {
@@ -39,6 +40,7 @@ describe('ReactFlightDOMBrowser', () => {
     ReactDOMServer = require('react-dom/server.browser');
     ReactServerDOMWriter = require('react-server-dom-webpack/writer.browser.server');
     ReactServerDOMReader = require('react-server-dom-webpack');
+    Suspense = React.Suspense;
   });
 
   function moduleReference(moduleExport) {
@@ -108,6 +110,11 @@ describe('ReactFlightDOMBrowser', () => {
     return [DelayedText, _resolve, _reject];
   }
 
+  const theInfinitePromise = new Promise(() => {});
+  function InfiniteSuspend() {
+    throw theInfinitePromise;
+  }
+
   it('should resolve HTML using W3C streams', async () => {
     function Text({children}) {
       return <span>{children}</span>;
@@ -180,7 +187,6 @@ describe('ReactFlightDOMBrowser', () => {
 
   it('should progressively reveal server components', async () => {
     let reportedErrors = [];
-    const {Suspense} = React;
 
     // Client Components
 
@@ -356,8 +362,6 @@ describe('ReactFlightDOMBrowser', () => {
   });
 
   it('should close the stream upon completion when rendering to W3C streams', async () => {
-    const {Suspense} = React;
-
     // Model
     function Text({children}) {
       return children;
@@ -512,4 +516,68 @@ describe('ReactFlightDOMBrowser', () => {
     const result = await readResult(ssrStream);
     expect(result).toEqual('<span>Client Component</span>');
   });
+
+  it('should be able to complete after aborting and throw the reason client-side', async () => {
+    const reportedErrors = [];
+
+    class ErrorBoundary extends React.Component {
+      state = {hasError: false, error: null};
+      static getDerivedStateFromError(error) {
+        return {
+          hasError: true,
+          error,
+        };
+      }
+      render() {
+        if (this.state.hasError) {
+          return this.props.fallback(this.state.error);
+        }
+        return this.props.children;
+      }
+    }
+
+    const controller = new AbortController();
+    const stream = ReactServerDOMWriter.renderToReadableStream(
+      <div>
+        <InfiniteSuspend />
+      </div>,
+      webpackMap,
+      {
+        signal: controller.signal,
+        onError(x) {
+          reportedErrors.push(x);
+        },
+      },
+    );
+    const response = ReactServerDOMReader.createFromReadableStream(stream);
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+
+    function App({res}) {
+      return res.readRoot();
+    }
+
+    await act(async () => {
+      root.render(
+        <ErrorBoundary fallback={e => <p>{e.message}</p>}>
+          <Suspense fallback={<p>(loading)</p>}>
+            <App res={response} />
+          </Suspense>
+        </ErrorBoundary>,
+      );
+    });
+    expect(container.innerHTML).toBe('<p>(loading)</p>');
+
+    await act(async () => {
+      // @TODO this is a hack to work around lack of support for abortSignal.reason in node
+      // The abort call itself should set this property but since we are testing in node we
+      // set it here manually
+      controller.signal.reason = 'for reasons';
+      controller.abort('for reasons');
+    });
+    expect(container.innerHTML).toBe('<p>Error: for reasons</p>');
+
+    expect(reportedErrors).toEqual(['for reasons']);
+  });
 });
diff --git a/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js b/packages/react-server-native-relay/src/ReactFlightNativeRelayServerHostConfig.js
@@ -114,6 +114,14 @@ export function processModelChunk(
   return ['J', id, json];
 }
 
+export function processReferenceChunk(
+  request: Request,
+  id: number,
+  reference: string,
+): Chunk {
+  return ['J', id, reference];
+}
+
 export function processModuleChunk(
   request: Request,
   id: number,
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
diff --git a/packages/react-server/src/ReactFlightServerConfigStream.js b/packages/react-server/src/ReactFlightServerConfigStream.js
