[Flight] Allow aborting encodeReply (#31106)

Allow aborting encoding arguments to a Server Action if a Promise doesn't resolve. That way at least part of the arguments can be used on the receiving side. This leaves it unresolved in the stream rather than encoding an error. This should error on the receiving side when the stream closes but it doesn't right now in the Edge/Browser versions because closing happens immediately before we've had a chance to call `.then()` so the Chunks are still in pending state. This is an existing bug also in FlightClient.
facebook · Oct 1, 2024 · 99c056a · 99c056a
1 parent d8c90fa
commit 99c056a
Show file tree

Hide file tree

Showing 7 changed files with 102 additions and 11 deletions.
diff --git a/packages/react-client/src/ReactFlightReplyClient.js b/packages/react-client/src/ReactFlightReplyClient.js
@@ -185,7 +185,7 @@ export function processReply(
   temporaryReferences: void | TemporaryReferenceSet,
   resolve: (string | FormData) => void,
   reject: (error: mixed) => void,
-): void {
+): (reason: mixed) => void {
   let nextPartId = 1;
   let pendingParts = 0;
   let formData: null | FormData = null;
@@ -841,6 +841,19 @@ export function processReply(
     return JSON.stringify(model, resolveToJSON);
   }
 
+  function abort(reason: mixed): void {
+    if (pendingParts > 0) {
+      pendingParts = 0; // Don't resolve again later.
+      // Resolve with what we have so far, which may have holes at this point.
+      // They'll error when the stream completes on the server.
+      if (formData === null) {
+        resolve(json);
+      } else {
+        resolve(formData);
+      }
+    }
+  }
+
   const json = serializeModel(root, 0);
 
   if (formData === null) {
@@ -854,6 +867,8 @@ export function processReply(
       resolve(formData);
     }
   }
+
+  return abort;
 }
 
 const boundCache: WeakMap<

diff --git a/packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js b/packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js
@@ -121,12 +121,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -135,6 +135,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

diff --git a/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientBrowser.js b/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientBrowser.js
@@ -120,12 +120,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -134,6 +134,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

diff --git a/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientEdge.js b/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientEdge.js
@@ -149,12 +149,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -163,6 +163,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
@@ -618,4 +618,20 @@ describe('ReactFlightDOMReply', () => {
     const root = await ReactServerDOMServer.decodeReply(body, webpackServerMap);
     expect(root.prop.obj).toBe(root.prop);
   });
+
+  it('can abort an unresolved model and get the partial result', async () => {
+    const promise = new Promise(r => {});
+    const controller = new AbortController();
+    const bodyPromise = ReactServerDOMClient.encodeReply(
+      {promise: promise, hello: 'world'},
+      {signal: controller.signal},
+    );
+    controller.abort();
+
+    const result = await ReactServerDOMServer.decodeReply(await bodyPromise);
+    expect(result.hello).toBe('world');
+    // TODO: await result.promise should reject at this point because the stream
+    // has closed but that's a bug in both ReactFlightReplyServer and ReactFlightClient.
+    // It just halts in this case.
+  });
 });
diff --git a/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientBrowser.js b/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientBrowser.js
@@ -120,12 +120,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -134,6 +134,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }
 

diff --git a/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientEdge.js b/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientEdge.js
@@ -149,12 +149,12 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
-  options?: {temporaryReferences?: TemporaryReferenceSet},
+  options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(
+    const abort = processReply(
       value,
       '',
       options && options.temporaryReferences
@@ -163,6 +163,18 @@ function encodeReply(
       resolve,
       reject,
     );
+    if (options && options.signal) {
+      const signal = options.signal;
+      if (signal.aborted) {
+        abort((signal: any).reason);
+      } else {
+        const listener = () => {
+          abort((signal: any).reason);
+          signal.removeEventListener('abort', listener);
+        };
+        signal.addEventListener('abort', listener);
+      }
+    }
   });
 }