Stop flowing and then abort if a stream is cancelled

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 () => {

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);
           },
         },

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);
           },
         },

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);
           },
         },

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;
     },

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},

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},

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);
     },
   };

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},

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},

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) {

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.',
+    ]);
+  });
 });

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 {

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 {