[Flight] Update stale blocked values in `createModelResolver` (#28669)
Alternative to #28620. Instead of emitting lazy references to not-yet-emitted models in the Flight Server, this fixes the observed issue in https://github.com/unstubbable/ai-rsc-test/pull/1 by adjusting the lazy model resolution in the Flight Client to update stale blocked root models, before assigning them as chunk values. In addition, the element props are not outlined anymore in the Flight Server to avoid having to also handle their staleness in blocked elements. fixes #28595
Hendrik Liebau committed
Apr 1, 2024 at 18:37 UTC
93f91795a0c71bae4aadd7f082b91de0068a0f91
3 files changed
+128
-13
packages/react-client/src/ReactFlightClient.js
+7
@@ -584,6 +584,13 @@ function createModelResolver<T>(
584
}
585
return value => {
586
parentObject[key] = value;
587
+
588
+ // If this is the root object for a model reference, where `blocked.value`
589
+ // is a stale `null`, the resolved value can be used directly.
590
+ if (key === '' && blocked.value === null) {
591
+ blocked.value = value;
592
+ }
593
+
594
blocked.deps--;
595
if (blocked.deps === 0) {
596
if (chunk.status !== BLOCKED) {
packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
+104
@@ -812,6 +812,110 @@ describe('ReactFlightDOM', () => {
812
expect(reportedErrors).toEqual([]);
813
});
814
815
+ it('should handle streaming async server components', async () => {
816
+ const reportedErrors = [];
817
+
818
+ const Row = async ({current, next}) => {
819
+ const chunk = await next;
820
+
821
+ if (chunk.done) {
822
+ return chunk.value;
823
+ }
824
+
825
+ return (
826
+ <Suspense fallback={chunk.value}>
827
+ <Row current={chunk.value} next={chunk.next} />
828
+ </Suspense>
829
+ );
830
+ };
831
+
832
+ function createResolvablePromise() {
833
+ let _resolve, _reject;
834
+
835
+ const promise = new Promise((resolve, reject) => {
836
+ _resolve = resolve;
837
+ _reject = reject;
838
+ });
839
+
840
+ return {promise, resolve: _resolve, reject: _reject};
841
+ }
842
+
843
+ function createSuspendedChunk(initialValue) {
844
+ const {promise, resolve, reject} = createResolvablePromise();
845
+
846
+ return {
847
+ row: (
848
+ <Suspense fallback={initialValue}>
849
+ <Row current={initialValue} next={promise} />
850
+ </Suspense>
851
+ ),
852
+ resolve,
853
+ reject,
854
+ };
855
+ }
856
+
857
+ function makeDelayedText() {
858
+ const {promise, resolve, reject} = createResolvablePromise();
859
+ async function DelayedText() {
860
+ const data = await promise;
861
+ return <div>{data}</div>;
862
+ }
863
+ return [DelayedText, resolve, reject];
864
+ }
865
+
866
+ const [Posts, resolvePostsData] = makeDelayedText();
867
+ const [Photos, resolvePhotosData] = makeDelayedText();
868
+ const suspendedChunk = createSuspendedChunk(<p>loading</p>);
869
+ const {writable, readable} = getTestStream();
870
+ const {pipe} = ReactServerDOMServer.renderToPipeableStream(
871
+ suspendedChunk.row,
872
+ webpackMap,
873
+ {
874
+ onError(error) {
875
+ reportedErrors.push(error);
876
+ },
877
+ },
878
+ );
879
+ pipe(writable);
880
+ const response = ReactServerDOMClient.createFromReadableStream(readable);
881
+ const container = document.createElement('div');
882
+ const root = ReactDOMClient.createRoot(container);
883
+
884
+ function ClientRoot() {
885
+ return use(response);
886
+ }
887
+
888
+ await act(() => {
889
+ root.render(<ClientRoot />);
890
+ });
891
+
892
+ expect(container.innerHTML).toBe('<p>loading</p>');
893
+
894
+ const donePromise = createResolvablePromise();
895
+
896
+ const value = (
897
+ <Suspense fallback={<p>loading posts and photos</p>}>
898
+ <Posts />
899
+ <Photos />
900
+ </Suspense>
901
+ );
902
+
903
+ await act(async () => {
904
+ suspendedChunk.resolve({value, done: false, next: donePromise.promise});
905
+ donePromise.resolve({value, done: true});
906
+ });
907
+
908
+ expect(container.innerHTML).toBe('<p>loading posts and photos</p>');
909
+
910
+ await act(async () => {
911
+ await resolvePostsData('posts');
912
+ await resolvePhotosData('photos');
913
+ });
914
+
915
+ expect(container.innerHTML).toBe('<div>posts</div><div>photos</div>');
916
+ expect(reportedErrors).toEqual([]);
917
+ });
918
+
919
it('should preserve state of client components on refetch', async () => {
920
// Client
921
packages/react-server/src/ReactFlightServer.js
+17
-13
@@ -243,11 +243,16 @@ export type ReactClientValue =
243
244
type ReactClientObject = {+[key: string]: ReactClientValue};
245
246
+// task status
247
const PENDING = 0;
248
const COMPLETED = 1;
249
const ABORTED = 3;
250
const ERRORED = 4;
251
252
+// object reference status
253
+const SEEN_BUT_NOT_YET_OUTLINED = -1;
254
+const NEVER_OUTLINED = -2;
255
+
256
type Task = {
257
id: number,
258
status: 0 | 1 | 3 | 4,
@@ -280,7 +285,7 @@ export type Request = {
285
writtenSymbols: Map<symbol, number>,
286
writtenClientReferences: Map<ClientReferenceKey, number>,
287
writtenServerReferences: Map<ServerReference<any>, number>,
283
- writtenObjects: WeakMap<Reference, number>, // -1 means "seen" but not outlined.
288
+ writtenObjects: WeakMap<Reference, number>,
289
identifierPrefix: string,
290
identifierCount: number,
291
taintCleanupQueue: Array<string | bigint>,
@@ -1125,8 +1130,7 @@ function serializeMap(
1130
const writtenObjects = request.writtenObjects;
1131
const existingId = writtenObjects.get(key);
1132
if (existingId === undefined) {
1128
- // Mark all object keys as seen so that they're always outlined.
1129
- writtenObjects.set(key, -1);
1133
+ writtenObjects.set(key, SEEN_BUT_NOT_YET_OUTLINED);
1134
}
1135
}
1136
}
@@ -1142,8 +1146,7 @@ function serializeSet(request: Request, set: Set<ReactClientValue>): string {
1146
const writtenObjects = request.writtenObjects;
1147
const existingId = writtenObjects.get(key);
1148
if (existingId === undefined) {
1145
- // Mark all object keys as seen so that they're always outlined.
1146
- writtenObjects.set(key, -1);
1149
+ writtenObjects.set(key, SEEN_BUT_NOT_YET_OUTLINED);
1150
}
1151
}
1152
}
@@ -1328,8 +1331,7 @@ function renderModelDestructive(
1331
// This is the ID we're currently emitting so we need to write it
1332
// once but if we discover it again, we refer to it by id.
1333
modelRoot = null;
1331
- } else if (existingId === -1) {
1332
- // Seen but not yet outlined.
1334
+ } else if (existingId === SEEN_BUT_NOT_YET_OUTLINED) {
1335
// TODO: If we throw here we can treat this as suspending which causes an outline
1336
// but that is able to reuse the same task if we're already in one but then that
1337
// will be a lazy future value rather than guaranteed to exist but maybe that's good.
@@ -1348,7 +1350,10 @@ function renderModelDestructive(
1350
} else {
1351
// This is the first time we've seen this object. We may never see it again
1352
// so we'll inline it. Mark it as seen. If we see it again, we'll outline.
1351
- writtenObjects.set(value, -1);
1353
+ writtenObjects.set(value, SEEN_BUT_NOT_YET_OUTLINED);
1354
+ // The element's props are marked as "never outlined" so that they are inlined into
1355
+ // the same row as the element itself.
1356
+ writtenObjects.set((value: any).props, NEVER_OUTLINED);
1357
}
1358
1359
const element: React$Element<any> = (value: any);
@@ -1477,11 +1482,10 @@ function renderModelDestructive(
1482
// This is the ID we're currently emitting so we need to write it
1483
// once but if we discover it again, we refer to it by id.
1484
modelRoot = null;
1480
- } else if (existingId === -1) {
1481
- // Seen but not yet outlined.
1485
+ } else if (existingId === SEEN_BUT_NOT_YET_OUTLINED) {
1486
const newId = outlineModel(request, (value: any));
1487
return serializeByValueID(newId);
1484
- } else {
1488
+ } else if (existingId !== NEVER_OUTLINED) {
1489
// We've already emitted this as an outlined object, so we can
1490
// just refer to that by its existing ID.
1491
return serializeByValueID(existingId);
@@ -1489,7 +1493,7 @@ function renderModelDestructive(
1493
} else {
1494
// This is the first time we've seen this object. We may never see it again
1495
// so we'll inline it. Mark it as seen. If we see it again, we'll outline.
1492
- writtenObjects.set(value, -1);
1496
+ writtenObjects.set(value, SEEN_BUT_NOT_YET_OUTLINED);
1497
}
1498
1499
if (isArray(value)) {
@@ -2007,7 +2011,7 @@ function renderConsoleValue(
2011
return serializeInfinitePromise();
2012
}
2013
2010
- if (existingId !== undefined && existingId !== -1) {
2014
+ if (existingId !== undefined && existingId >= 0) {
2015
// We've already emitted this as a real object, so we can
2016
// just refer to that by its existing ID.
2017
return serializeByValueID(existingId);