Serialize byte stream chunks as 'b' rows instead of o rows.

This allows us to avoid using the `_pendingByteStreamIDs` Set in the
client to track which streams are byte streams.

Author: unstubbable

packages/react-client/src/ReactFlightClient.js

@@ -2758,7 +2758,6 @@ export type StreamState = {
   _rowTag: number, // 0 indicates that we're currently parsing the row ID
   _rowLength: number, // remaining bytes in the row. 0 indicates that we're looking for a newline.
   _buffer: Array<Uint8Array>, // chunks received so far as part of this row
-  _pendingByteStreamIDs: Set<number>, // IDs of pending streams with type: 'bytes'
   _debugInfo: ReactIOInfo, // DEV-only
   _debugTargetChunkSize: number, // DEV-only
 };
@@ -2773,7 +2772,6 @@ export function createStreamState(
     _rowTag: 0,
     _rowLength: 0,
     _buffer: [],
-    _pendingByteStreamIDs: new Set(),
   }: Omit<StreamState, '_debugInfo' | '_debugTargetChunkSize'>): any);
   if (__DEV__ && enableAsyncDebugInfo) {
     const response = unwrapWeakResponse(weakResponse);
@@ -4790,7 +4788,6 @@ function processFullStringRow(
     }
     // Fallthrough
     case 114 /* "r" */: {
-      streamState._pendingByteStreamIDs.add(id);
       startReadableStream(response, id, 'bytes', streamState);
       return;
     }
@@ -4806,9 +4803,6 @@ function processFullStringRow(
     }
     // Fallthrough
     case 67 /* "C" */: {
-      if (streamState._pendingByteStreamIDs.has(id)) {
-        streamState._pendingByteStreamIDs.delete(id);
-      }
       stopStream(response, id, row);
       return;
     }
@@ -4864,6 +4858,7 @@ export function processBinaryChunk(
           resolvedRowTag === 65 /* "A" */ ||
           resolvedRowTag === 79 /* "O" */ ||
           resolvedRowTag === 111 /* "o" */ ||
+          resolvedRowTag === 98 /* "b" */ ||
           resolvedRowTag === 85 /* "U" */ ||
           resolvedRowTag === 83 /* "S" */ ||
           resolvedRowTag === 115 /* "s" */ ||
@@ -4926,10 +4921,7 @@ export function processBinaryChunk(
 
       // Check if this is a Uint8Array for a byte stream. We enqueue it
       // immediately but need to determine if we can use zero-copy or must copy.
-      if (
-        rowTag === 111 /* "o" */ &&
-        streamState._pendingByteStreamIDs.has(rowID)
-      ) {
+      if (rowTag === 98 /* "b" */) {
         resolveBuffer(
           response,
           rowID,
@@ -4970,10 +4962,7 @@ export function processBinaryChunk(
       // For byte streams, we can enqueue the partial chunk immediately using
       // zero-copy since we're at the end of the RSC chunk and no more parsing
       // will access this buffer.
-      if (
-        rowTag === 111 /* "o" */ &&
-        streamState._pendingByteStreamIDs.has(rowID)
-      ) {
+      if (rowTag === 98 /* "b" */) {
         // Update how many bytes we're still waiting for. We need to do this
         // before enqueueing as enqueue will detach the buffer and byteLength
         // will become 0.

packages/react-server/src/ReactFlightServer.js

@@ -1149,6 +1149,8 @@ function serializeReadableStream(
       supportsBYOB = false;
     }
   }
+  // At this point supportsBYOB is guaranteed to be a boolean.
+  const isByteStream: boolean = supportsBYOB;
 
   const reader = stream.getReader();
 
@@ -1172,7 +1174,7 @@ function serializeReadableStream(
   // The task represents the Stop row. This adds a Start row.
   request.pendingChunks++;
   const startStreamRow =
-    streamTask.id.toString(16) + ':' + (supportsBYOB ? 'r' : 'R') + '\n';
+    streamTask.id.toString(16) + ':' + (isByteStream ? 'r' : 'R') + '\n';
   request.completedRegularChunks.push(stringToChunk(startStreamRow));
 
   function progress(entry: {done: boolean, value: ReactClientValue, ...}) {
@@ -1192,7 +1194,7 @@ function serializeReadableStream(
       try {
         streamTask.model = entry.value;
         request.pendingChunks++;
-        tryStreamTask(request, streamTask);
+        tryStreamTask(request, streamTask, isByteStream);
         enqueueFlush(request);
         reader.read().then(progress, error);
       } catch (x) {
@@ -1319,7 +1321,7 @@ function serializeAsyncIterable(
       try {
         streamTask.model = entry.value;
         request.pendingChunks++;
-        tryStreamTask(request, streamTask);
+        tryStreamTask(request, streamTask, false);
         enqueueFlush(request);
         if (__DEV__) {
           callIteratorInDEV(iterator, progress, error);
@@ -5534,6 +5536,7 @@ function emitChunk(
   request: Request,
   task: Task,
   value: ReactClientValue,
+  isByteStream: boolean,
 ): void {
   const id = task.id;
   // For certain types we have special types, we typically outlined them but
@@ -5559,7 +5562,7 @@ function emitChunk(
   }
   if (value instanceof Uint8Array) {
     // unsigned char
-    emitTypedArrayChunk(request, id, 'o', value, false);
+    emitTypedArrayChunk(request, id, isByteStream ? 'b' : 'o', value, false);
     return;
   }
   if (value instanceof Uint8ClampedArray) {
@@ -5717,7 +5720,7 @@ function retryTask(request: Request, task: Task): void {
 
       // Object might contain unresolved values like additional elements.
       // This is simulating what the JSON loop would do if this was part of it.
-      emitChunk(request, task, resolvedModel);
+      emitChunk(request, task, resolvedModel, false);
     } 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.
@@ -5776,7 +5779,11 @@ function retryTask(request: Request, task: Task): void {
   }
 }
 
-function tryStreamTask(request: Request, task: Task): void {
+function tryStreamTask(
+  request: Request,
+  task: Task,
+  isByteStream: boolean,
+): void {
   // This is used to try to emit something synchronously but if it suspends,
   // we emit a reference to a new outlined task immediately instead.
   const prevCanEmitDebugInfo = canEmitDebugInfo;
@@ -5787,7 +5794,7 @@ function tryStreamTask(request: Request, task: Task): void {
   }
   const parentSerializedSize = serializedSize;
   try {
-    emitChunk(request, task, task.model);
+    emitChunk(request, task, task.model, isByteStream);
   } finally {
     serializedSize = parentSerializedSize;
     if (__DEV__) {