[Fizz] Expose a method to abort a pending request (facebook#21027)

sebmarkbage · zhengjitf · commit 9c5faecabed7 · 2021-03-23T11:40:18.000+08:00
* Track all suspended work while it's still pending

This allows us to abort work and put everything into client rendered mode
if we don't want to wait for further I/O.

It also allows us to cancel fallbacks if we complete the main content
before the fallback.

* Expose abort API to the browser streams

Since this API already returns a value, we need to use destructuring to
expose more options.

* Add a test including the client actually client rendering it

* Use AbortSignal option for W3C streams instead of external control

* Clean up listener after it's used once
diff --git a/fixtures/fizz-ssr-browser/index.html b/fixtures/fizz-ssr-browser/index.html
@@ -20,7 +20,10 @@ <h1>Fizz Example</h1>
     <script src="../../build/node_modules/react-dom/umd/react-dom-unstable-fizz.browser.development.js"></script>
     <script src="https://unpkg.com/babel-standalone@6/babel.js"></script>
     <script type="text/babel">
-      let stream = ReactDOMFizzServer.renderToReadableStream(<body>Success</body>);
+      let controller = new AbortController();
+      let stream = ReactDOMFizzServer.renderToReadableStream(<body>Success</body>, {
+        signal: controller.signal,
+      });
       let response = new Response(stream, {
         headers: {'Content-Type': 'text/html'},
       });
diff --git a/package.json b/package.json
@@ -36,6 +36,7 @@
     "@babel/preset-react": "^7.10.4",
     "@babel/traverse": "^7.11.0",
     "@mattiasbuelens/web-streams-polyfill": "^0.3.2",
+    "abort-controller": "^3.0.0",
     "art": "0.10.1",
     "babel-eslint": "^10.0.3",
     "babel-plugin-syntax-trailing-function-commas": "^6.5.0",
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -263,4 +263,55 @@ describe('ReactDOMFizzServer', () => {
     // Now it's hydrated.
     expect(ref.current).toBe(h1);
   });
+
+  // @gate experimental
+  it('client renders a boundary if it does not resolve before aborting', async () => {
+    function App() {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <h1>
+              <AsyncText text="Hello" />
+            </h1>
+          </Suspense>
+        </div>
+      );
+    }
+
+    let controls;
+    await act(async () => {
+      controls = ReactDOMFizzServer.pipeToNodeWritable(<App />, writable);
+    });
+
+    // We're still showing a fallback.
+
+    // Attempt to hydrate the content.
+    const root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+    Scheduler.unstable_flushAll();
+
+    // We're still loading because we're waiting for the server to stream more content.
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    // We abort the server response.
+    await act(async () => {
+      controls.abort();
+    });
+
+    // We still can't render it on the client.
+    Scheduler.unstable_flushAll();
+    expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
+
+    // We now resolve it on the client.
+    resolveText('Hello');
+
+    Scheduler.unstable_flushAll();
+
+    // The client rendered HTML is now in place.
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <h1>Hello</h1>
+      </div>,
+    );
+  });
 });
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
@@ -12,6 +12,7 @@
 // Polyfills for test environment
 global.ReadableStream = require('@mattiasbuelens/web-streams-polyfill/ponyfill/es6').ReadableStream;
 global.TextEncoder = require('util').TextEncoder;
+global.AbortController = require('abort-controller');
 
 let React;
 let ReactDOMFizzServer;
@@ -110,4 +111,22 @@ describe('ReactDOMFizzServer', () => {
     const result = await readResult(stream);
     expect(result).toContain('Loading');
   });
+
+  // @gate experimental
+  it('should be able to complete by aborting even if the promise never resolves', async () => {
+    const controller = new AbortController();
+    const stream = ReactDOMFizzServer.renderToReadableStream(
+      <div>
+        <Suspense fallback={<div>Loading</div>}>
+          <InfiniteSuspend />
+        </Suspense>
+      </div>,
+      {signal: controller.signal},
+    );
+
+    controller.abort();
+
+    const result = await readResult(stream);
+    expect(result).toContain('Loading');
+  });
 });
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerNode-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerNode-test.js
@@ -115,4 +115,54 @@ describe('ReactDOMFizzServer', () => {
     expect(output.error).toBe(undefined);
     expect(output.result).toContain('Loading');
   });
+
+  // @gate experimental
+  it('should not attempt to render the fallback if the main content completes first', async () => {
+    const {writable, output, completed} = getTestWritable();
+
+    let renderedFallback = false;
+    function Fallback() {
+      renderedFallback = true;
+      return 'Loading...';
+    }
+    function Content() {
+      return 'Hi';
+    }
+    ReactDOMFizzServer.pipeToNodeWritable(
+      <Suspense fallback={<Fallback />}>
+        <Content />
+      </Suspense>,
+      writable,
+    );
+
+    await completed;
+
+    expect(output.result).toContain('Hi');
+    expect(output.result).not.toContain('Loading');
+    expect(renderedFallback).toBe(false);
+  });
+
+  // @gate experimental
+  it('should be able to complete by aborting even if the promise never resolves', async () => {
+    const {writable, output, completed} = getTestWritable();
+    const {abort} = ReactDOMFizzServer.pipeToNodeWritable(
+      <div>
+        <Suspense fallback={<div>Loading</div>}>
+          <InfiniteSuspend />
+        </Suspense>
+      </div>,
+      writable,
+    );
+
+    jest.runAllTimers();
+
+    expect(output.result).toContain('Loading');
+
+    abort();
+
+    await completed;
+
+    expect(output.error).toBe(undefined);
+    expect(output.result).toContain('Loading');
+  });
 });
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
@@ -13,10 +13,26 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  abort,
 } from 'react-server/src/ReactFizzServer';
 
-function renderToReadableStream(children: ReactNodeList): ReadableStream {
+type Options = {
+  signal?: AbortSignal,
+};
+
+function renderToReadableStream(
+  children: ReactNodeList,
+  options?: Options,
+): ReadableStream {
   let request;
+  if (options && options.signal) {
+    const signal = options.signal;
+    const listener = () => {
+      abort(request);
+      signal.removeEventListener('abort', listener);
+    };
+    signal.addEventListener('abort', listener);
+  }
   return new ReadableStream({
     start(controller) {
       request = createRequest(children, controller);
diff --git a/packages/react-dom/src/server/ReactDOMFizzServerNode.js b/packages/react-dom/src/server/ReactDOMFizzServerNode.js
@@ -14,19 +14,31 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  abort,
 } from 'react-server/src/ReactFizzServer';
 
 function createDrainHandler(destination, request) {
   return () => startFlowing(request);
 }
 
+type Controls = {
+  // Cancel any pending I/O and put anything remaining into
+  // client rendered mode.
+  abort(): void,
+};
+
 function pipeToNodeWritable(
   children: ReactNodeList,
   destination: Writable,
-): void {
+): Controls {
   const request = createRequest(children, destination);
   destination.on('drain', createDrainHandler(destination, request));
   startWork(request);
+  return {
+    abort() {
+      abort(request);
+    },
+  };
 }
 
 export {pipeToNodeWritable};
diff --git a/packages/react-noop-renderer/src/ReactNoopServer.js b/packages/react-noop-renderer/src/ReactNoopServer.js
@@ -216,6 +216,9 @@ function render(children: React$Element<any>): Destination {
     placeholders: new Map(),
     segments: new Map(),
     stack: [],
+    abort() {
+      ReactNoopServer.abort(request);
+    },
   };
   const request = ReactNoopServer.createRequest(children, destination);
   ReactNoopServer.startWork(request);
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
