From 4fd4f935ee31114ae3bdf6b57fc56c6be3bed212 Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Fri, 8 Apr 2022 14:34:37 -0700
Subject: [PATCH] Implement encodeInto during flushCompletedQueues

encodeInto allows us to write directly to the view buffer that will end up getting streamed instead of encoding into an intermediate buffer and then copying that data.
---
 .../src/ReactServerStreamConfigNode.js        | 66 +++++++++++++++++--
 scripts/flow/environment.js                   | 27 ++++++++
 2 files changed, 86 insertions(+), 7 deletions(-)

diff --git a/packages/react-server/src/ReactServerStreamConfigNode.js b/packages/react-server/src/ReactServerStreamConfigNode.js
index 02f529bbeb29..c34d288cd2a4 100644
--- a/packages/react-server/src/ReactServerStreamConfigNode.js
+++ b/packages/react-server/src/ReactServerStreamConfigNode.js
@@ -18,7 +18,7 @@ type MightBeFlushable = {
 export type Destination = Writable & MightBeFlushable;
 
 export type PrecomputedChunk = Uint8Array;
-export type Chunk = Uint8Array;
+export type Chunk = string;
 
 export function scheduleWork(callback: () => void) {
   setImmediate(callback);
@@ -45,14 +45,49 @@ export function beginWriting(destination: Destination) {
   destinationHasCapacity = true;
 }
 
-export function writeChunk(
-  destination: Destination,
-  chunk: PrecomputedChunk | Chunk,
-): void {
-  if (chunk.byteLength === 0) {
+function writeStringChunk(destination: Destination, stringChunk: string) {
+  if (stringChunk.length === 0) {
+    return;
+  }
+  // maximum possible view needed to encode entire string
+  if (stringChunk.length * 3 > VIEW_SIZE) {
+    if (writtenBytes > 0) {
+      writeToDestination(
+        destination,
+        ((currentView: any): Uint8Array).subarray(0, writtenBytes),
+      );
+      currentView = new Uint8Array(VIEW_SIZE);
+      writtenBytes = 0;
+    }
+    writeToDestination(destination, textEncoder.encode(stringChunk));
     return;
   }
 
+  let target: Uint8Array = (currentView: any);
+  if (writtenBytes > 0) {
+    target = ((currentView: any): Uint8Array).subarray(writtenBytes);
+  }
+  const {read, written} = textEncoder.encodeInto(stringChunk, target);
+  writtenBytes += written;
+
+  if (read < stringChunk.length) {
+    writeToDestination(destination, (currentView: any));
+    currentView = new Uint8Array(VIEW_SIZE);
+    writtenBytes = textEncoder.encodeInto(stringChunk.slice(read), currentView)
+      .written;
+  }
+
+  if (writtenBytes === VIEW_SIZE) {
+    writeToDestination(destination, (currentView: any));
+    currentView = new Uint8Array(VIEW_SIZE);
+    writtenBytes = 0;
+  }
+}
+
+function writeViewChunk(destination: Destination, chunk: PrecomputedChunk) {
+  if (chunk.byteLength === 0) {
+    return;
+  }
   if (chunk.byteLength > VIEW_SIZE) {
     // this chunk may overflow a single view which implies it was not
     // one that is cached by the streaming renderer. We will enqueu
@@ -93,6 +128,23 @@ export function writeChunk(
   }
   ((currentView: any): Uint8Array).set(bytesToWrite, writtenBytes);
   writtenBytes += bytesToWrite.byteLength;
+
+  if (writtenBytes === VIEW_SIZE) {
+    writeToDestination(destination, (currentView: any));
+    currentView = new Uint8Array(VIEW_SIZE);
+    writtenBytes = 0;
+  }
+}
+
+export function writeChunk(
+  destination: Destination,
+  chunk: PrecomputedChunk | Chunk,
+): void {
+  if (typeof chunk === 'string') {
+    writeStringChunk(destination, chunk);
+  } else {
+    writeViewChunk(destination, ((chunk: any): PrecomputedChunk));
+  }
 }
 
 function writeToDestination(destination: Destination, view: Uint8Array) {
@@ -124,7 +176,7 @@ export function close(destination: Destination) {
 const textEncoder = new TextEncoder();
 
 export function stringToChunk(content: string): Chunk {
-  return textEncoder.encode(content);
+  return content;
 }
 
 export function stringToPrecomputedChunk(content: string): PrecomputedChunk {
diff --git a/scripts/flow/environment.js b/scripts/flow/environment.js
index 4294964a74b9..f97bd32a4321 100644
--- a/scripts/flow/environment.js
+++ b/scripts/flow/environment.js
@@ -125,6 +125,33 @@ declare module 'pg' {
   };
 }
 
+declare module 'util' {
+  declare function debuglog(section: string): (data: any, ...args: any) => void;
+  declare function format(format: string, ...placeholders: any): string;
+  declare function log(string: string): void;
+  declare function inspect(object: any, options?: util$InspectOptions): string;
+  declare function isArray(object: any): boolean;
+  declare function isRegExp(object: any): boolean;
+  declare function isDate(object: any): boolean;
+  declare function isError(object: any): boolean;
+  declare function inherits(
+    constructor: Function,
+    superConstructor: Function,
+  ): void;
+  declare function deprecate(f: Function, string: string): Function;
+  declare function promisify(f: Function): Function;
+  declare function callbackify(f: Function): Function;
+  declare class TextEncoder {
+    constructor(encoding?: string): TextEncoder;
+    encode(buffer: string): Uint8Array;
+    encodeInto(
+      buffer: string,
+      dest: Uint8Array,
+    ): {read: number, written: number};
+    encoding: string;
+  }
+}
+
 declare module 'pg/lib/utils' {
   declare module.exports: {
     prepareValue(val: any): mixed,