[Flight Reply] Reject any new Chunks not yet discovered at the time of reportGlobalError (#31840)

sebmarkbage · web-flow · commit ef979d4703ff · 2024-12-18T12:56:57.000-08:00
We might have already resolved models that are not pending and so are
not rejected by aborting the stream. When those later get parsed they
might discover new chunks which end up as pending. These should be
errored since they will never be able to resolve later.

This avoids infinitely hanging the stream.

This same fix needs to be ported to ReactFlightClient that has the same
issue.
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js
@@ -245,4 +245,31 @@ describe('ReactFlightDOMReplyEdge', () => {
       ),
     );
   });
+
+  it('should abort when parsing an incomplete payload', async () => {
+    const infinitePromise = new Promise(() => {});
+    const controller = new AbortController();
+    const promiseForResult = ReactServerDOMClient.encodeReply(
+      {promise: infinitePromise},
+      {
+        signal: controller.signal,
+      },
+    );
+    controller.abort();
+    const body = await promiseForResult;
+
+    const decoded = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+
+    let error = null;
+    try {
+      await decoded.promise;
+    } catch (x) {
+      error = x;
+    }
+    expect(error).not.toBe(null);
+    expect(error.message).toBe('Connection closed.');
+  });
 });
diff --git a/packages/react-server/src/ReactFlightReplyServer.js b/packages/react-server/src/ReactFlightReplyServer.js
@@ -169,6 +169,8 @@ export type Response = {
   _prefix: string,
   _formData: FormData,
   _chunks: Map<number, SomeChunk<any>>,
+  _closed: boolean,
+  _closedReason: mixed,
   _temporaryReferences: void | TemporaryReferenceSet,
 };
 
@@ -255,6 +257,14 @@ function createResolvedModelChunk<T>(
   return new Chunk(RESOLVED_MODEL, value, id, response);
 }
 
+function createErroredChunk<T>(
+  response: Response,
+  reason: mixed,
+): ErroredChunk<T> {
+  // $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
+  return new Chunk(ERRORED, null, reason, response);
+}
+
 function resolveModelChunk<T>(
   chunk: SomeChunk<T>,
   value: string,
@@ -493,6 +503,8 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
 // Report that any missing chunks in the model is now going to throw this
 // error upon read. Also notify any pending promises.
 export function reportGlobalError(response: Response, error: Error): void {
+  response._closed = true;
+  response._closedReason = error;
   response._chunks.forEach(chunk => {
     // If this chunk was already resolved or errored, it won't
     // trigger an error but if it wasn't then we need to
@@ -514,6 +526,10 @@ function getChunk(response: Response, id: number): SomeChunk<any> {
     if (backingEntry != null) {
       // We assume that this is a string entry for now.
       chunk = createResolvedModelChunk(response, (backingEntry: any), id);
+    } else if (response._closed) {
+      // We have already errored the response and we're not going to get
+      // anything more streaming in so this will immediately error.
+      chunk = createErroredChunk(response, response._closedReason);
     } else {
       // We're still waiting on this entry to stream in.
       chunk = createPendingChunk(response);
@@ -1082,6 +1098,8 @@ export function createResponse(
     _prefix: formFieldPrefix,
     _formData: backingFormData,
     _chunks: chunks,
+    _closed: false,
+    _closedReason: null,
     _temporaryReferences: temporaryReferences,
   };
   return response;
