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.
facebook · Apr 8, 2022 · d036841 · d036841
1 parent fc976d7
commit d036841
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 9 deletions.
diff --git 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
@@ -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,

diff --git a/scripts/rollup/bundles.js b/scripts/rollup/bundles.js
@@ -318,7 +318,7 @@ const bundles = [
     global: 'ReactDOMServer',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
-    externals: ['react'],
+    externals: ['react', 'util'],
   },
   {
     bundleTypes: __EXPERIMENTAL__ ? [FB_WWW_DEV, FB_WWW_PROD] : [],
@@ -347,7 +347,7 @@ const bundles = [
     global: 'ReactServerDOMWriter',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
-    externals: ['react'],
+    externals: ['react', 'util'],
   },
 
   /******* React Server DOM Webpack Reader *******/
@@ -437,6 +437,7 @@ const bundles = [
       'ReactFlightNativeRelayServerIntegration',
       'JSResourceReferenceImpl',
       'ReactNativeInternalFeatureFlags',
+      'util',
     ],
   },