[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.
Sebastian Markbåge committed
Oct 1, 2024 at 11:28 UTC
99c056abb0dac0e1a15b2c85b620b72c625e065b
7 files changed
+102
-11
packages/react-client/src/ReactFlightReplyClient.js
+16
-1
@@ -185,7 +185,7 @@ export function processReply(
185
temporaryReferences: void | TemporaryReferenceSet,
186
resolve: (string | FormData) => void,
187
reject: (error: mixed) => void,
188
-): void {
188
+): (reason: mixed) => void {
189
let nextPartId = 1;
190
let pendingParts = 0;
191
let formData: null | FormData = null;
@@ -841,6 +841,19 @@ export function processReply(
841
return JSON.stringify(model, resolveToJSON);
842
}
843
844
+ function abort(reason: mixed): void {
845
+ if (pendingParts > 0) {
846
+ pendingParts = 0; // Don't resolve again later.
847
+ // Resolve with what we have so far, which may have holes at this point.
848
+ // They'll error when the stream completes on the server.
849
+ if (formData === null) {
850
+ resolve(json);
851
+ } else {
852
+ resolve(formData);
853
+ }
854
+ }
855
+ }
856
+
857
const json = serializeModel(root, 0);
858
859
if (formData === null) {
@@ -854,6 +867,8 @@ export function processReply(
867
resolve(formData);
868
}
869
}
870
+
871
+ return abort;
872
}
873
874
const boundCache: WeakMap<
packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js
+14
-2
@@ -121,12 +121,12 @@ function createFromFetch<T>(
121
122
function encodeReply(
123
value: ReactServerValue,
124
- options?: {temporaryReferences?: TemporaryReferenceSet},
124
+ options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
125
): Promise<
126
string | URLSearchParams | FormData,
127
> /* We don't use URLSearchParams yet but maybe */ {
128
return new Promise((resolve, reject) => {
129
- processReply(
129
+ const abort = processReply(
130
value,
131
'',
132
options && options.temporaryReferences
@@ -135,6 +135,18 @@ function encodeReply(
135
resolve,
136
reject,
137
);
138
+ if (options && options.signal) {
139
+ const signal = options.signal;
140
+ if (signal.aborted) {
141
+ abort((signal: any).reason);
142
+ } else {
143
+ const listener = () => {
144
+ abort((signal: any).reason);
145
+ signal.removeEventListener('abort', listener);
146
+ };
147
+ signal.addEventListener('abort', listener);
148
+ }
149
+ }
150
});
151
}
152
packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientBrowser.js
+14
-2
@@ -120,12 +120,12 @@ function createFromFetch<T>(
120
121
function encodeReply(
122
value: ReactServerValue,
123
- options?: {temporaryReferences?: TemporaryReferenceSet},
123
+ options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
124
): Promise<
125
string | URLSearchParams | FormData,
126
> /* We don't use URLSearchParams yet but maybe */ {
127
return new Promise((resolve, reject) => {
128
- processReply(
128
+ const abort = processReply(
129
value,
130
'',
131
options && options.temporaryReferences
@@ -134,6 +134,18 @@ function encodeReply(
134
resolve,
135
reject,
136
);
137
+ if (options && options.signal) {
138
+ const signal = options.signal;
139
+ if (signal.aborted) {
140
+ abort((signal: any).reason);
141
+ } else {
142
+ const listener = () => {
143
+ abort((signal: any).reason);
144
+ signal.removeEventListener('abort', listener);
145
+ };
146
+ signal.addEventListener('abort', listener);
147
+ }
148
+ }
149
});
150
}
151
packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientEdge.js
+14
-2
@@ -149,12 +149,12 @@ function createFromFetch<T>(
149
150
function encodeReply(
151
value: ReactServerValue,
152
- options?: {temporaryReferences?: TemporaryReferenceSet},
152
+ options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
153
): Promise<
154
string | URLSearchParams | FormData,
155
> /* We don't use URLSearchParams yet but maybe */ {
156
return new Promise((resolve, reject) => {
157
- processReply(
157
+ const abort = processReply(
158
value,
159
'',
160
options && options.temporaryReferences
@@ -163,6 +163,18 @@ function encodeReply(
163
resolve,
164
reject,
165
);
166
+ if (options && options.signal) {
167
+ const signal = options.signal;
168
+ if (signal.aborted) {
169
+ abort((signal: any).reason);
170
+ } else {
171
+ const listener = () => {
172
+ abort((signal: any).reason);
173
+ signal.removeEventListener('abort', listener);
174
+ };
175
+ signal.addEventListener('abort', listener);
176
+ }
177
+ }
178
});
179
}
180
packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
+16
@@ -618,4 +618,20 @@ describe('ReactFlightDOMReply', () => {
618
const root = await ReactServerDOMServer.decodeReply(body, webpackServerMap);
619
expect(root.prop.obj).toBe(root.prop);
620
});
621
+
622
+ it('can abort an unresolved model and get the partial result', async () => {
623
+ const promise = new Promise(r => {});
624
+ const controller = new AbortController();
625
+ const bodyPromise = ReactServerDOMClient.encodeReply(
626
+ {promise: promise, hello: 'world'},
627
+ {signal: controller.signal},
628
+ );
629
+ controller.abort();
630
+
631
+ const result = await ReactServerDOMServer.decodeReply(await bodyPromise);
632
+ expect(result.hello).toBe('world');
633
+ // TODO: await result.promise should reject at this point because the stream
634
+ // has closed but that's a bug in both ReactFlightReplyServer and ReactFlightClient.
635
+ // It just halts in this case.
636
+ });
637
});
packages/react-server-dom-webpack/src/client/ReactFlightDOMClientBrowser.js
+14
-2
@@ -120,12 +120,12 @@ function createFromFetch<T>(
120
121
function encodeReply(
122
value: ReactServerValue,
123
- options?: {temporaryReferences?: TemporaryReferenceSet},
123
+ options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
124
): Promise<
125
string | URLSearchParams | FormData,
126
> /* We don't use URLSearchParams yet but maybe */ {
127
return new Promise((resolve, reject) => {
128
- processReply(
128
+ const abort = processReply(
129
value,
130
'',
131
options && options.temporaryReferences
@@ -134,6 +134,18 @@ function encodeReply(
134
resolve,
135
reject,
136
);
137
+ if (options && options.signal) {
138
+ const signal = options.signal;
139
+ if (signal.aborted) {
140
+ abort((signal: any).reason);
141
+ } else {
142
+ const listener = () => {
143
+ abort((signal: any).reason);
144
+ signal.removeEventListener('abort', listener);
145
+ };
146
+ signal.addEventListener('abort', listener);
147
+ }
148
+ }
149
});
150
}
151
packages/react-server-dom-webpack/src/client/ReactFlightDOMClientEdge.js
+14
-2
@@ -149,12 +149,12 @@ function createFromFetch<T>(
149
150
function encodeReply(
151
value: ReactServerValue,
152
- options?: {temporaryReferences?: TemporaryReferenceSet},
152
+ options?: {temporaryReferences?: TemporaryReferenceSet, signal?: AbortSignal},
153
): Promise<
154
string | URLSearchParams | FormData,
155
> /* We don't use URLSearchParams yet but maybe */ {
156
return new Promise((resolve, reject) => {
157
- processReply(
157
+ const abort = processReply(
158
value,
159
'',
160
options && options.temporaryReferences
@@ -163,6 +163,18 @@ function encodeReply(
163
resolve,
164
reject,
165
);
166
+ if (options && options.signal) {
167
+ const signal = options.signal;
168
+ if (signal.aborted) {
169
+ abort((signal: any).reason);
170
+ } else {
171
+ const listener = () => {
172
+ abort((signal: any).reason);
173
+ signal.removeEventListener('abort', listener);
174
+ };
175
+ signal.addEventListener('abort', listener);
176
+ }
177
+ }
178
});
179
}
180