Skip to content

Commit

Permalink
Emit root strings or typed arrays without outlining
Browse files Browse the repository at this point in the history
  • Loading branch information
@sebmarkbage
sebmarkbage committed Apr 16, 2024
1 parent cecd166 commit e8d5dea
Show file tree
Hide file tree
Showing 4 changed files with 207 additions and 45 deletions.
76 changes: 69 additions & 7 deletions packages/react-client/src/ReactFlightClient.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,8 @@ import {
export type {CallServerCallback, EncodeFormActionCallback};

interface FlightStreamController {
enqueue(json: UninitializedModel): void;
enqueueValue(value: any): void;
enqueueModel(json: UninitializedModel): void;
close(json: UninitializedModel): void;
error(error: Error): void;
}
Expand Down Expand Up @@ -381,6 +382,15 @@ function createInitializedBufferChunk(
return new Chunk(INITIALIZED, value, null, response);
}

function createInitializedIteratorResultChunk<T>(
response: Response,
value: T,
done: boolean,
): InitializedChunk<IteratorResult<T, T>> {
// $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
return new Chunk(INITIALIZED, {done: done, value: value}, null, response);
}

function createInitializedStreamChunk<
T: ReadableStream | $AsyncIterable<any, any, void>,
>(
Expand Down Expand Up @@ -427,7 +437,7 @@ function resolveModelChunk<T>(
// a stream chunk since any other row shouldn't have more than one entry.
const streamChunk: InitializedStreamChunk<any> = (chunk: any);
const controller = streamChunk.reason;
controller.enqueue(value);
controller.enqueueModel(value);
}
return;
}
Expand Down Expand Up @@ -1034,8 +1044,17 @@ function resolveModel(

function resolveText(response: Response, id: number, text: string): void {
const chunks = response._chunks;
// We assume that we always reference large strings after they've been
// emitted.
if (enableFlightReadableStream) {
const chunk = chunks.get(id);
if (chunk && chunk.status !== PENDING) {
// If we get more data to an already resolved ID, we assume that it's
// a stream chunk since any other row shouldn't have more than one entry.
const streamChunk: InitializedStreamChunk<any> = (chunk: any);
const controller = streamChunk.reason;
controller.enqueueValue(text);
return;
}
}
chunks.set(id, createInitializedTextChunk(response, text));
}

Expand All @@ -1045,7 +1064,17 @@ function resolveBuffer(
buffer: $ArrayBufferView | ArrayBuffer,
): void {
const chunks = response._chunks;
// We assume that we always reference buffers after they've been emitted.
if (enableFlightReadableStream) {
const chunk = chunks.get(id);
if (chunk && chunk.status !== PENDING) {
// If we get more data to an already resolved ID, we assume that it's
// a stream chunk since any other row shouldn't have more than one entry.
const streamChunk: InitializedStreamChunk<any> = (chunk: any);
const controller = streamChunk.reason;
controller.enqueueValue(buffer);
return;
}
}
chunks.set(id, createInitializedBufferChunk(response, buffer));
}

Expand Down Expand Up @@ -1143,7 +1172,17 @@ function startReadableStream<T>(
});
let previousBlockedChunk: SomeChunk<T> | null = null;
const flightController = {
enqueue(json: UninitializedModel): void {
enqueueValue(value: T): void {
if (previousBlockedChunk === null) {
controller.enqueue(value);
} else {
// We're still waiting on a previous chunk so we can't enqueue quite yet.
previousBlockedChunk.then(function () {
controller.enqueue(value);
});
}
},
enqueueModel(json: UninitializedModel): void {
if (previousBlockedChunk === null) {
// If we're not blocked on any other chunks, we can try to eagerly initialize
// this as a fast-path to avoid awaiting them.
Expand Down Expand Up @@ -1236,7 +1275,30 @@ function startAsyncIterable<T>(
let closed = false;
let nextWriteIndex = 0;
const flightController = {
enqueue(value: UninitializedModel): void {
enqueueValue(value: T): void {
if (nextWriteIndex === buffer.length) {
buffer[nextWriteIndex] = createInitializedIteratorResultChunk(
response,
value,
false,
);
} else {
const chunk: PendingChunk<IteratorResult<T, T>> = (buffer[
nextWriteIndex
]: any);
const resolveListeners = chunk.value;
const rejectListeners = chunk.reason;
const initializedChunk: InitializedChunk<IteratorResult<T, T>> =
(chunk: any);
initializedChunk.status = INITIALIZED;
initializedChunk.value = {done: false, value: value};
if (resolveListeners !== null) {
wakeChunkIfInitialized(chunk, resolveListeners, rejectListeners);
}
}
nextWriteIndex++;
},
enqueueModel(value: UninitializedModel): void {
if (nextWriteIndex === buffer.length) {
buffer[nextWriteIndex] = createResolvedIteratorResultChunk(
response,
Expand Down
4 changes: 2 additions & 2 deletions packages/react-client/src/__tests__/ReactFlight-test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2305,7 +2305,6 @@ describe('ReactFlight', () => {
return {
async *[Symbol.asyncIterator]() {
yield <span>Who</span>;
yield ' ';
yield <span>dis?</span>;
resolve();
},
Expand Down Expand Up @@ -2386,7 +2385,8 @@ describe('ReactFlight', () => {

expect(ReactNoop).toMatchRenderedOutput(
<div>
<span>Who</span> <span>dis?</span>
<span>Who</span>
<span>dis?</span>
</div>,
);
});
Expand Down
4 changes: 2 additions & 2 deletions packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -615,7 +615,7 @@ describe('ReactFlightDOMEdge', () => {
},
);

expect(await readByteLength(stream2)).toBeLessThan(400);
expect(await readByteLength(stream2)).toBeLessThan(300);

const streamedBuffers = [];
const reader = result.getReader();
Expand Down Expand Up @@ -672,7 +672,7 @@ describe('ReactFlightDOMEdge', () => {
},
);

expect(await readByteLength(stream2)).toBeLessThan(400);
expect(await readByteLength(stream2)).toBeLessThan(300);

const streamedBuffers = [];
const reader = result.getReader({mode: 'byob'});
Expand Down
168 changes: 134 additions & 34 deletions packages/react-server/src/ReactFlightServer.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1411,13 +1411,9 @@ function serializeTemporaryReference(
}

function serializeLargeTextString(request: Request, text: string): string {
request.pendingChunks += 2;
request.pendingChunks++;
const textId = request.nextChunkId++;
const textChunk = stringToChunk(text);
const binaryLength = byteLengthOfChunk(textChunk);
const row = textId.toString(16) + ':T' + binaryLength.toString(16) + ',';
const headerChunk = stringToChunk(row);
request.completedRegularChunks.push(headerChunk, textChunk);
emitTextChunk(request, textId, text);
return serializeByValueID(textId);
}

Expand Down Expand Up @@ -1467,27 +1463,9 @@ function serializeTypedArray(
tag: string,
typedArray: $ArrayBufferView,
): string {
if (enableTaint) {
if (TaintRegistryByteLengths.has(typedArray.byteLength)) {
// If we have had any tainted values of this length, we check
// to see if these bytes matches any entries in the registry.
const tainted = TaintRegistryValues.get(
binaryToComparableString(typedArray),
);
if (tainted !== undefined) {
throwTaintViolation(tainted.message);
}
}
}
request.pendingChunks += 2;
request.pendingChunks++;
const bufferId = request.nextChunkId++;
// TODO: Convert to little endian if that's not the server default.
const binaryChunk = typedArrayToBinaryChunk(typedArray);
const binaryLength = byteLengthOfBinaryChunk(binaryChunk);
const row =
bufferId.toString(16) + ':' + tag + binaryLength.toString(16) + ',';
const headerChunk = stringToChunk(row);
request.completedRegularChunks.push(headerChunk, binaryChunk);
emitTypedArrayChunk(request, bufferId, tag, typedArray);
return serializeByValueID(bufferId);
}

Expand Down Expand Up @@ -2321,6 +2299,42 @@ function emitDebugChunk(
request.completedRegularChunks.push(processedChunk);
}

function emitTypedArrayChunk(
request: Request,
id: number,
tag: string,
typedArray: $ArrayBufferView,
): void {
if (enableTaint) {
if (TaintRegistryByteLengths.has(typedArray.byteLength)) {
// If we have had any tainted values of this length, we check
// to see if these bytes matches any entries in the registry.
const tainted = TaintRegistryValues.get(
binaryToComparableString(typedArray),
);
if (tainted !== undefined) {
throwTaintViolation(tainted.message);
}
}
}
request.pendingChunks++; // Extra chunk for the header.
// TODO: Convert to little endian if that's not the server default.
const binaryChunk = typedArrayToBinaryChunk(typedArray);
const binaryLength = byteLengthOfBinaryChunk(binaryChunk);
const row = id.toString(16) + ':' + tag + binaryLength.toString(16) + ',';
const headerChunk = stringToChunk(row);
request.completedRegularChunks.push(headerChunk, binaryChunk);
}

function emitTextChunk(request: Request, id: number, text: string): void {
request.pendingChunks++; // Extra chunk for the header.
const textChunk = stringToChunk(text);
const binaryLength = byteLengthOfChunk(textChunk);
const row = id.toString(16) + ':T' + binaryLength.toString(16) + ',';
const headerChunk = stringToChunk(row);
request.completedRegularChunks.push(headerChunk, textChunk);
}

function serializeEval(source: string): string {
if (!__DEV__) {
// These errors should never make it into a build so we don't need to encode them in codes.json
Expand Down Expand Up @@ -2681,6 +2695,96 @@ function forwardDebugInfo(
}
}

function emitChunk(
request: Request,
task: Task,
value: ReactClientValue,
): void {
const id = task.id;
// For certain types we have special types, we typically outlined them but
// we can emit them directly for this row instead of through an indirection.
if (typeof value === 'string') {
if (enableTaint) {
const tainted = TaintRegistryValues.get(value);
if (tainted !== undefined) {
throwTaintViolation(tainted.message);
}
}
emitTextChunk(request, id, value);
return;
}
if (enableBinaryFlight) {
if (value instanceof ArrayBuffer) {
emitTypedArrayChunk(request, id, 'A', new Uint8Array(value));
return;
}
if (value instanceof Int8Array) {
// char
emitTypedArrayChunk(request, id, 'O', value);
return;
}
if (value instanceof Uint8Array) {
// unsigned char
emitTypedArrayChunk(request, id, 'o', value);
return;
}
if (value instanceof Uint8ClampedArray) {
// unsigned clamped char
emitTypedArrayChunk(request, id, 'U', value);
return;
}
if (value instanceof Int16Array) {
// sort
emitTypedArrayChunk(request, id, 'S', value);
return;
}
if (value instanceof Uint16Array) {
// unsigned short
emitTypedArrayChunk(request, id, 's', value);
return;
}
if (value instanceof Int32Array) {
// long
emitTypedArrayChunk(request, id, 'L', value);
return;
}
if (value instanceof Uint32Array) {
// unsigned long
emitTypedArrayChunk(request, id, 'l', value);
return;
}
if (value instanceof Float32Array) {
// float
emitTypedArrayChunk(request, id, 'G', value);
return;
}
if (value instanceof Float64Array) {
// double
emitTypedArrayChunk(request, id, 'g', value);
return;
}
if (value instanceof BigInt64Array) {
// number
emitTypedArrayChunk(request, id, 'M', value);
return;
}
if (value instanceof BigUint64Array) {
// unsigned number
// We use "m" instead of "n" since JSON can start with "null"
emitTypedArrayChunk(request, id, 'm', value);
return;
}
if (value instanceof DataView) {
emitTypedArrayChunk(request, id, 'V', value);
return;
}
}
// For anything else we need to try to serialize it using JSON.
// $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
const json: string = stringify(value, task.toJSON);
emitModelChunk(request, task.id, json);
}

const emptyRoot = {};

function retryTask(request: Request, task: Task): void {
Expand Down Expand Up @@ -2725,19 +2829,17 @@ function retryTask(request: Request, task: Task): void {
task.keyPath = null;
task.implicitSlot = false;

let json: string;
if (typeof resolvedModel === 'object' && resolvedModel !== null) {
// Object might contain unresolved values like additional elements.
// This is simulating what the JSON loop would do if this was part of it.
// $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
json = stringify(resolvedModel, task.toJSON);
emitChunk(request, task, resolvedModel);
} else {
// If the value is a string, it means it's a terminal value and we already escaped it
// We don't need to escape it again so it's not passed the toJSON replacer.
// $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
json = stringify(resolvedModel);
const json: string = stringify(resolvedModel);
emitModelChunk(request, task.id, json);
}
emitModelChunk(request, task.id, json);

request.abortableTasks.delete(task);
task.status = COMPLETED;
Expand Down Expand Up @@ -2789,9 +2891,7 @@ function tryStreamTask(request: Request, task: Task): void {
debugID = null;
}
try {
// $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
const json: string = stringify(task.model, task.toJSON);
emitModelChunk(request, task.id, json);
emitChunk(request, task, task.model);
} finally {
if (__DEV__) {
debugID = prevDebugID;
Expand Down

0 comments on commit e8d5dea

Please sign in to comment.