Enable bundle reader for Node. (#3167)

Author: wu-hui

packages/firestore/src/platform/platform.ts

@@ -18,8 +18,9 @@
 import { DatabaseId, DatabaseInfo } from '../core/database_info';
 import { Connection } from '../remote/connection';
 import { JsonProtoSerializer } from '../remote/serializer';
-import { fail } from '../util/assert';
+import { debugAssert, fail } from '../util/assert';
 import { ConnectivityMonitor } from './../remote/connectivity_monitor';
+import { BundleSource } from '../util/bundle_reader';
 
 /**
  * Provides a common interface to load anything platform dependent, e.g.
@@ -50,6 +51,18 @@ export interface Platform {
    */
   randomBytes(nBytes: number): Uint8Array;
 
+  /**
+   * Builds a `ByteStreamReader` from a data source.
+   * @param source The data source to use.
+   * @param bytesPerRead How many bytes each `read()` from the returned reader
+   *        will read. It is ignored if the passed in source does not provide
+   *        such control(example: ReadableStream).
+   */
+  toByteStreamReader(
+    source: BundleSource,
+    bytesPerRead: number
+  ): ReadableStreamReader<Uint8Array>;
+
   /** The Platform's 'window' implementation or null if not available. */
   readonly window: Window | null;
 
@@ -60,6 +73,40 @@ export interface Platform {
   readonly base64Available: boolean;
 }
 
+/**
+ * Builds a `ByteStreamReader` from a UInt8Array.
+ * @param source The data source to use.
+ * @param bytesPerRead How many bytes each `read()` from the returned reader
+ *        will read.
+ */
+export function toByteStreamReader(
+  source: Uint8Array,
+  bytesPerRead: number
+): ReadableStreamReader<Uint8Array> {
+  debugAssert(
+    bytesPerRead > 0,
+    `toByteStreamReader expects positive bytesPerRead, but got ${bytesPerRead}`
+  );
+  let readFrom = 0;
+  const reader: ReadableStreamReader<Uint8Array> = {
+    async read(): Promise<ReadableStreamReadResult<Uint8Array>> {
+      if (readFrom < source.byteLength) {
+        const result = {
+          value: source.slice(readFrom, readFrom + bytesPerRead),
+          done: false
+        };
+        readFrom += bytesPerRead;
+        return result;
+      }
+
+      return { value: undefined, done: true };
+    },
+    async cancel(): Promise<void> {},
+    releaseLock() {}
+  };
+  return reader;
+}
+
 /**
  * Provides singleton helpers where setup code can inject a platform at runtime.
  * setPlatform needs to be set before Firestore is used and must be set exactly

packages/firestore/src/platform_browser/browser_platform.ts

@@ -16,7 +16,7 @@
  */
 
 import { DatabaseId, DatabaseInfo } from '../core/database_info';
-import { Platform } from '../platform/platform';
+import { Platform, toByteStreamReader } from '../platform/platform';
 import { Connection } from '../remote/connection';
 import { JsonProtoSerializer } from '../remote/serializer';
 import { ConnectivityMonitor } from './../remote/connectivity_monitor';
@@ -25,6 +25,7 @@ import { NoopConnectivityMonitor } from '../remote/connectivity_monitor_noop';
 import { BrowserConnectivityMonitor } from './browser_connectivity_monitor';
 import { WebChannelConnection } from './webchannel_connection';
 import { debugAssert } from '../util/assert';
+import { BundleSource } from '../util/bundle_reader';
 
 // Implements the Platform API for browsers and some browser-like environments
 // (including ReactNative).
@@ -93,4 +94,25 @@ export class BrowserPlatform implements Platform {
     }
     return bytes;
   }
+
+  /**
+   * On web, a `ReadableStream` is wrapped around by a `ByteStreamReader`.
+   */
+  toByteStreamReader(
+    source: BundleSource,
+    bytesPerRead: number
+  ): ReadableStreamReader<Uint8Array> {
+    if (source instanceof Uint8Array) {
+      return toByteStreamReader(source, bytesPerRead);
+    }
+    if (source instanceof ArrayBuffer) {
+      return toByteStreamReader(new Uint8Array(source), bytesPerRead);
+    }
+    if (source instanceof ReadableStream) {
+      return source.getReader();
+    }
+    throw new Error(
+      'Source of `toByteStreamReader` has to be a ArrayBuffer or ReadableStream'
+    );
+  }
 }

packages/firestore/src/platform_node/node_platform.ts

@@ -19,7 +19,7 @@ import { randomBytes } from 'crypto';
 import { inspect } from 'util';
 
 import { DatabaseId, DatabaseInfo } from '../core/database_info';
-import { Platform } from '../platform/platform';
+import { Platform, toByteStreamReader } from '../platform/platform';
 import { Connection } from '../remote/connection';
 import { JsonProtoSerializer } from '../remote/serializer';
 import { Code, FirestoreError } from '../util/error';
@@ -29,6 +29,7 @@ import { NoopConnectivityMonitor } from './../remote/connectivity_monitor_noop';
 import { GrpcConnection } from './grpc_connection';
 import { loadProtos } from './load_protos';
 import { debugAssert } from '../util/assert';
+import { invalidClassError } from '../util/input_validation';
 
 export class NodePlatform implements Platform {
   readonly base64Available = true;
@@ -83,4 +84,22 @@ export class NodePlatform implements Platform {
 
     return randomBytes(nBytes);
   }
+
+  /**
+   * On Node, only supported data source is a `Uint8Array` for now.
+   */
+  toByteStreamReader(
+    source: Uint8Array,
+    bytesPerRead: number
+  ): ReadableStreamReader<Uint8Array> {
+    if (!(source instanceof Uint8Array)) {
+      throw invalidClassError(
+        'NodePlatform.toByteStreamReader',
+        'Uint8Array',
+        1,
+        source
+      );
+    }
+    return toByteStreamReader(source, bytesPerRead);
+  }
 }

packages/firestore/src/util/bundle_reader.ts

@@ -20,6 +20,8 @@ import {
   BundleMetadata
 } from '../protos/firestore_bundle_proto';
 import { Deferred } from './promise';
+import { PlatformSupport } from '../platform/platform';
+import { debugAssert } from './assert';
 
 /**
  * A complete element in the bundle stream, together with the byte length it
@@ -37,29 +39,18 @@ export class SizedBundleElement {
   }
 }
 
+export type BundleSource =
+  | ReadableStream<Uint8Array>
+  | ArrayBuffer
+  | Uint8Array;
+
 /**
- * Create a `ReadableStream` from a underlying buffer.
+ * When applicable, how many bytes to read from the underlying data source
+ * each time.
  *
- * @param data: Underlying buffer.
- * @param bytesPerRead: How many bytes to read from the underlying buffer from
- * each read through the stream.
+ * Not applicable for ReadableStreams.
  */
-export function toReadableStream(
-  data: Uint8Array | ArrayBuffer,
-  bytesPerRead = 10240
-): ReadableStream<Uint8Array | ArrayBuffer> {
-  let readFrom = 0;
-  return new ReadableStream({
-    start(controller) {},
-    async pull(controller): Promise<void> {
-      controller.enqueue(data.slice(readFrom, readFrom + bytesPerRead));
-      readFrom += bytesPerRead;
-      if (readFrom >= data.byteLength) {
-        controller.close();
-      }
-    }
-  });
-}
+const BYTES_PER_READ = 10240;
 
 /**
  * A class representing a bundle.
@@ -70,8 +61,6 @@ export function toReadableStream(
 export class BundleReader {
   /** Cached bundle metadata. */
   private metadata: Deferred<BundleMetadata> = new Deferred<BundleMetadata>();
-  /** The reader instance of the given ReadableStream. */
-  private reader: ReadableStreamDefaultReader;
   /**
    * Internal buffer to hold bundle content, accumulating incomplete element
    * content.
@@ -80,20 +69,16 @@ export class BundleReader {
   /** The decoder used to parse binary data into strings. */
   private textDecoder = new TextDecoder('utf-8');
 
+  static fromBundleSource(source: BundleSource): BundleReader {
+    return new BundleReader(
+      PlatformSupport.getPlatform().toByteStreamReader(source, BYTES_PER_READ)
+    );
+  }
+
   constructor(
-    private bundleStream:
-      | ReadableStream<Uint8Array | ArrayBuffer>
-      | Uint8Array
-      | ArrayBuffer
+    /** The reader to read from underlying binary bundle data source. */
+    private reader: ReadableStreamReader<Uint8Array>
   ) {
-    if (
-      bundleStream instanceof Uint8Array ||
-      bundleStream instanceof ArrayBuffer
-    ) {
-      this.bundleStream = toReadableStream(bundleStream);
-    }
-    this.reader = (this.bundleStream as ReadableStream).getReader();
-
     // Read the metadata (which is the first element).
     this.nextElementImpl().then(
       element => {
@@ -220,8 +205,8 @@ export class BundleReader {
 
   private raiseError(message: string): void {
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
-    this.reader.cancel('Invalid bundle format.');
-    throw new Error(message);
+    this.reader.cancel();
+    throw new Error(`Invalid bundle format: ${message}`);
   }
 
   /**
@@ -231,11 +216,12 @@ export class BundleReader {
   private async pullMoreDataToBuffer(): Promise<boolean> {
     const result = await this.reader.read();
     if (!result.done) {
+      debugAssert(!!result.value, 'Read undefined when "done" is false.');
       const newBuffer = new Uint8Array(
-        this.buffer.length + result.value.length
+        this.buffer.length + result.value!.length
       );
       newBuffer.set(this.buffer);
-      newBuffer.set(result.value, this.buffer.length);
+      newBuffer.set(result.value!, this.buffer.length);
       this.buffer = newBuffer;
     }
     return result.done;

packages/firestore/test/unit/platform/platform.test.ts

@@ -16,10 +16,37 @@
  */
 
 import { expect } from 'chai';
-import { PlatformSupport } from '../../../src/platform/platform';
+import {
+  PlatformSupport,
+  toByteStreamReader
+} from '../../../src/platform/platform';
 
 describe('Platform', () => {
   it('can load the platform at runtime', () => {
     expect(PlatformSupport.getPlatform()).to.exist;
   });
+
+  it('toByteStreamReader() steps underlying data', async () => {
+    const encoder = new TextEncoder();
+    const r = toByteStreamReader(
+      encoder.encode('0123456789'),
+      /* bytesPerRead */ 4
+    );
+
+    let result = await r.read();
+    expect(result.value).to.deep.equal(encoder.encode('0123'));
+    expect(result.done).to.be.false;
+
+    result = await r.read();
+    expect(result.value).to.deep.equal(encoder.encode('4567'));
+    expect(result.done).to.be.false;
+
+    result = await r.read();
+    expect(result.value).to.deep.equal(encoder.encode('89'));
+    expect(result.done).to.be.false;
+
+    result = await r.read();
+    expect(result.value).to.be.undefined;
+    expect(result.done).to.be.true;
+  });
 });