Stop flowing and then abort if a stream is cancelled (facebook#27405)

We currently abort a stream either it's explicitly told to abort (e.g. by an abortsignal). In this case we still finish writing what we have as well as instructions for the client about what happened so it can trigger fallback cases and log appropriately. We also abort a request if the stream itself cancels. E.g. if you can't write anymore. In this case we should not write anything to the outgoing stream since it's supposed to be closed already now. However, we should still abort the request so that more work isn't performed and so that we can log the reason for it to the onError callback. We should also not do any work after aborting. There we need to stop the "flow" of bytes - so I call stopFlowing in the cancel case before aborting. The tests were testing this case but we had changed the implementation to only start flowing at initial read (pull) instead of start like we used to. As a result, it was no longer covering this case. We have to call reader.read() in the tests to start the flow so that we need to cancel it. We also were missing a final assertion on the error logs and since we were tracking them explicitly the extra error was silenced.
EdisonVan · Apr 15, 2024 · 5ef1485 · 5ef1485
1 parent 3a81c97
commit 5ef1485
Show file tree

Hide file tree

Showing 14 changed files with 125 additions and 10 deletions.
diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServerBrowser-test.js
@@ -342,7 +342,8 @@ describe('ReactDOMFizzServerBrowser', () => {
     expect(isComplete).toBe(false);
 
     const reader = stream.getReader();
-    reader.cancel();
+    await reader.read();
+    await reader.cancel();
 
     expect(errors).toEqual([
       'The render was aborted by the server without a reason.',
@@ -355,6 +356,10 @@ describe('ReactDOMFizzServerBrowser', () => {
 
     expect(rendered).toBe(false);
     expect(isComplete).toBe(true);
+
+    expect(errors).toEqual([
+      'The render was aborted by the server without a reason.',
+    ]);
   });
 
   it('should stream large contents that might overlow individual buffers', async () => {

diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js b/packages/react-dom/src/server/ReactDOMFizzServerBrowser.js
@@ -19,6 +19,7 @@ import {
   resumeRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -78,6 +79,7 @@ function renderToReadableStream(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
@@ -158,6 +160,7 @@ function resume(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },

diff --git a/packages/react-dom/src/server/ReactDOMFizzServerBun.js b/packages/react-dom/src/server/ReactDOMFizzServerBun.js
@@ -17,6 +17,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -68,6 +69,7 @@ function renderToReadableStream(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },

diff --git a/packages/react-dom/src/server/ReactDOMFizzServerEdge.js b/packages/react-dom/src/server/ReactDOMFizzServerEdge.js
@@ -19,6 +19,7 @@ import {
   resumeRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -78,6 +79,7 @@ function renderToReadableStream(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },
@@ -158,6 +160,7 @@ function resume(
             startFlowing(request, controller);
           },
           cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
             abort(request);
           },
         },

diff --git a/packages/react-dom/src/server/ReactDOMFizzServerNode.js b/packages/react-dom/src/server/ReactDOMFizzServerNode.js
@@ -21,6 +21,7 @@ import {
   resumeRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFizzServer';
 
@@ -35,9 +36,12 @@ function createDrainHandler(destination: Destination, request: Request) {
   return () => startFlowing(request, destination);
 }
 
-function createAbortHandler(request: Request, reason: string) {
-  // eslint-disable-next-line react-internal/prod-error-codes
-  return () => abort(request, new Error(reason));
+function createCancelHandler(request: Request, reason: string) {
+  return () => {
+    stopFlowing(request);
+    // eslint-disable-next-line react-internal/prod-error-codes
+    abort(request, new Error(reason));
+  };
 }
 
 type Options = {
@@ -122,14 +126,14 @@ function renderToPipeableStream(
       destination.on('drain', createDrainHandler(destination, request));
       destination.on(
         'error',
-        createAbortHandler(
+        createCancelHandler(
           request,
           'The destination stream errored while writing data.',
         ),
       );
       destination.on(
         'close',
-        createAbortHandler(request, 'The destination stream closed early.'),
+        createCancelHandler(request, 'The destination stream closed early.'),
       );
       return destination;
     },
@@ -180,14 +184,14 @@ function resumeToPipeableStream(
       destination.on('drain', createDrainHandler(destination, request));
       destination.on(
         'error',
-        createAbortHandler(
+        createCancelHandler(
           request,
           'The destination stream errored while writing data.',
         ),
       );
       destination.on(
         'close',
-        createAbortHandler(request, 'The destination stream closed early.'),
+        createCancelHandler(request, 'The destination stream closed early.'),
       );
       return destination;
     },

diff --git a/packages/react-dom/src/server/ReactDOMFizzStaticBrowser.js b/packages/react-dom/src/server/ReactDOMFizzStaticBrowser.js
@@ -18,6 +18,7 @@ import {
   createPrerenderRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
   getPostponedState,
 } from 'react-server/src/ReactFizzServer';
@@ -61,6 +62,10 @@ function prerender(
           pull: (controller): ?Promise<void> => {
             startFlowing(request, controller);
           },
+          cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
+            abort(request);
+          },
         },
         // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
         {highWaterMark: 0},

diff --git a/packages/react-dom/src/server/ReactDOMFizzStaticEdge.js b/packages/react-dom/src/server/ReactDOMFizzStaticEdge.js
@@ -18,6 +18,7 @@ import {
   createPrerenderRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
   getPostponedState,
 } from 'react-server/src/ReactFizzServer';
@@ -61,6 +62,10 @@ function prerender(
           pull: (controller): ?Promise<void> => {
             startFlowing(request, controller);
           },
+          cancel: (reason): ?Promise<void> => {
+            stopFlowing(request);
+            abort(request);
+          },
         },
         // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
         {highWaterMark: 0},

diff --git a/packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js b/packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js
@@ -22,6 +22,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -90,6 +91,7 @@ function renderToPipeableStream(
       return destination;
     },
     abort(reason: mixed) {
+      stopFlowing(request);
       abort(request, reason);
     },
   };

diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js
@@ -16,6 +16,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -78,7 +79,10 @@ function renderToReadableStream(
       pull: (controller): ?Promise<void> => {
         startFlowing(request, controller);
       },
-      cancel: (reason): ?Promise<void> => {},
+      cancel: (reason): ?Promise<void> => {
+        stopFlowing(request);
+        abort(request, reason);
+      },
     },
     // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
     {highWaterMark: 0},

diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerEdge.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerEdge.js
@@ -16,6 +16,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -78,7 +79,10 @@ function renderToReadableStream(
       pull: (controller): ?Promise<void> => {
         startFlowing(request, controller);
       },
-      cancel: (reason): ?Promise<void> => {},
+      cancel: (reason): ?Promise<void> => {
+        stopFlowing(request);
+        abort(request, reason);
+      },
     },
     // $FlowFixMe[prop-missing] size() methods are not allowed on byte streams.
     {highWaterMark: 0},

diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js b/packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js
@@ -22,6 +22,7 @@ import {
   createRequest,
   startWork,
   startFlowing,
+  stopFlowing,
   abort,
 } from 'react-server/src/ReactFlightServer';
 
@@ -51,6 +52,14 @@ function createDrainHandler(destination: Destination, request: Request) {
   return () => startFlowing(request, destination);
 }
 
+function createCancelHandler(request: Request, reason: string) {
+  return () => {
+    stopFlowing(request);
+    // eslint-disable-next-line react-internal/prod-error-codes
+    abort(request, new Error(reason));
+  };
+}
+
 type Options = {
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
@@ -88,6 +97,17 @@ function renderToPipeableStream(
       hasStartedFlowing = true;
       startFlowing(request, destination);
       destination.on('drain', createDrainHandler(destination, request));
+      destination.on(
+        'error',
+        createCancelHandler(
+          request,
+          'The destination stream errored while writing data.',
+        ),
+      );
+      destination.on(
+        'close',
+        createCancelHandler(request, 'The destination stream closed early.'),
+      );
       return destination;
     },
     abort(reason: mixed) {

diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js
@@ -1222,4 +1222,53 @@ describe('ReactFlightDOMBrowser', () => {
 
     expect(postponed).toBe('testing postpone');
   });
+
+  it('should not continue rendering after the reader cancels', async () => {
+    let hasLoaded = false;
+    let resolve;
+    let rendered = false;
+    const promise = new Promise(r => (resolve = r));
+    function Wait() {
+      if (!hasLoaded) {
+        throw promise;
+      }
+      rendered = true;
+      return 'Done';
+    }
+    const errors = [];
+    const stream = await ReactServerDOMServer.renderToReadableStream(
+      <div>
+        <Suspense fallback={<div>Loading</div>}>
+          <Wait />
+        </Suspense>
+      </div>,
+      null,
+      {
+        onError(x) {
+          errors.push(x.message);
+        },
+      },
+    );
+
+    expect(rendered).toBe(false);
+
+    const reader = stream.getReader();
+    await reader.read();
+    await reader.cancel();
+
+    expect(errors).toEqual([
+      'The render was aborted by the server without a reason.',
+    ]);
+
+    hasLoaded = true;
+    resolve();
+
+    await jest.runAllTimers();
+
+    expect(rendered).toBe(false);
+
+    expect(errors).toEqual([
+      'The render was aborted by the server without a reason.',
+    ]);
+  });
 });
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
@@ -3998,6 +3998,10 @@ export function startFlowing(request: Request, destination: Destination): void {
   }
 }
 
+export function stopFlowing(request: Request): void {
+  request.destination = null;
+}
+
 // This is called to early terminate a request. It puts all pending boundaries in client rendered state.
 export function abort(request: Request, reason: mixed): void {
   try {

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
@@ -359,6 +359,7 @@ function serializeThenable(request: Request, thenable: Thenable<any>): number {
     },
     reason => {
       newTask.status = ERRORED;
+      request.abortableTasks.delete(newTask);
       // TODO: We should ideally do this inside performWork so it's scheduled
       const digest = logRecoverableError(request, reason);
       emitErrorChunk(request, newTask.id, digest, reason);
@@ -1570,6 +1571,10 @@ export function startFlowing(request: Request, destination: Destination): void {
   }
 }
 
+export function stopFlowing(request: Request): void {
+  request.destination = null;
+}
+
 // This is called to early terminate a request. It creates an error at all pending tasks.
 export function abort(request: Request, reason: mixed): void {
   try {