Support bundle reader for Node

Author: wu-hui

packages/firestore/src/platform/platform.ts

@@ -50,6 +50,8 @@ export interface Platform {
    */
   randomBytes(nBytes: number): Uint8Array;
 
+  toByteStreamReader(source: unknown): ByteStreamReader;
+
   /** The Platform's 'window' implementation or null if not available. */
   readonly window: Window | null;
 
@@ -60,6 +62,36 @@ export interface Platform {
   readonly base64Available: boolean;
 }
 
+export interface ByteStreamReader {
+  read(): Promise<ByteStreamReadResult>;
+}
+
+export interface ByteStreamReadResult {
+  done: boolean;
+  value: Uint8Array;
+}
+
+export function toByteStreamReader(
+  source: Uint8Array,
+  bytesPerRead = 10240
+): ByteStreamReader {
+  let readFrom = 0;
+  return new (class implements ByteStreamReader {
+    async read(): Promise<ByteStreamReadResult> {
+      if (readFrom < source.byteLength) {
+        const result = {
+          value: source.slice(readFrom, readFrom + bytesPerRead),
+          done: false
+        };
+        readFrom += bytesPerRead;
+        return result;
+      }
+
+      return { value: new Uint8Array(), done: true };
+    }
+  })();
+}
+
 /**
  * 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,12 @@
  */
 
 import { DatabaseId, DatabaseInfo } from '../core/database_info';
-import { Platform } from '../platform/platform';
+import {
+  ByteStreamReader,
+  ByteStreamReadResult,
+  Platform,
+  toByteStreamReader
+} from '../platform/platform';
 import { Connection } from '../remote/connection';
 import { JsonProtoSerializer } from '../remote/serializer';
 import { ConnectivityMonitor } from './../remote/connectivity_monitor';
@@ -93,4 +98,24 @@ export class BrowserPlatform implements Platform {
     }
     return bytes;
   }
+
+  toByteStreamReader(source: unknown): ByteStreamReader {
+    if (source instanceof Uint8Array) {
+      return toByteStreamReader(source);
+    }
+    if (source instanceof ArrayBuffer) {
+      return toByteStreamReader(new Uint8Array(source));
+    }
+    if (source instanceof ReadableStream) {
+      const reader = source.getReader();
+      return new (class implements ByteStreamReader {
+        read(): Promise<ByteStreamReadResult> {
+          return reader.read();
+        }
+      })();
+    }
+    throw new Error(
+      'Source of `toByteStreamReader` has to be Uint8Array, ArrayBuffer or ReadableStream'
+    );
+  }
 }

packages/firestore/src/platform_node/node_platform.ts

@@ -19,7 +19,11 @@ import { randomBytes } from 'crypto';
 import { inspect } from 'util';
 
 import { DatabaseId, DatabaseInfo } from '../core/database_info';
-import { Platform } from '../platform/platform';
+import {
+  ByteStreamReader,
+  Platform,
+  toByteStreamReader
+} from '../platform/platform';
 import { Connection } from '../remote/connection';
 import { JsonProtoSerializer } from '../remote/serializer';
 import { Code, FirestoreError } from '../util/error';
@@ -83,4 +87,11 @@ export class NodePlatform implements Platform {
 
     return randomBytes(nBytes);
   }
+
+  toByteStreamReader(source: unknown): ByteStreamReader {
+    if (source instanceof Uint8Array) {
+      return toByteStreamReader(source);
+    }
+    throw new Error('Source of `toByteStreamReader` has to be Uint8Array');
+  }
 }

packages/firestore/src/util/bundle_reader.ts

@@ -19,6 +19,7 @@ import {
   BundleElement,
   BundleMetadata
 } from '../protos/firestore_bundle_proto';
+import { ByteStreamReader, PlatformSupport } from '../platform/platform';
 
 /**
  * A complete element in the bundle stream, together with the byte length it
@@ -36,29 +37,6 @@ export class SizedBundleElement {
   }
 }
 
-/**
- * Create a `ReadableStream` from a underlying buffer.
- *
- * @param data: Underlying buffer.
- * @param bytesPerRead: How many bytes to read from the underlying buffer from each read through the stream.
- */
-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();
-      }
-    }
-  });
-}
-
 /**
  * A class representing a bundle.
  *
@@ -68,8 +46,8 @@ export function toReadableStream(
 export class BundleReader {
   // Cached bundle metadata.
   private metadata?: BundleMetadata | null;
-  // The reader instance of the given ReadableStream.
-  private reader: ReadableStreamDefaultReader;
+  // The reader to read from underlying binary bundle data source.
+  private reader: ByteStreamReader;
   // Internal buffer to hold bundle content, accumulating incomplete element content.
   private buffer: Uint8Array = new Uint8Array();
   private textDecoder = new TextDecoder('utf-8');
@@ -80,13 +58,9 @@ export class BundleReader {
       | Uint8Array
       | ArrayBuffer
   ) {
-    if (
-      bundleStream instanceof Uint8Array ||
-      bundleStream instanceof ArrayBuffer
-    ) {
-      this.bundleStream = toReadableStream(bundleStream);
-    }
-    this.reader = (this.bundleStream as ReadableStream).getReader();
+    this.reader = PlatformSupport.getPlatform().toByteStreamReader(
+      bundleStream
+    );
   }
 
   /**

packages/firestore/test/util/test_platform.ts

@@ -16,7 +16,7 @@
  */
 
 import { DatabaseId, DatabaseInfo } from '../../src/core/database_info';
-import { Platform } from '../../src/platform/platform';
+import { ByteStreamReader, Platform } from '../../src/platform/platform';
 import { Connection } from '../../src/remote/connection';
 import { JsonProtoSerializer } from '../../src/remote/serializer';
 import { debugAssert, fail } from '../../src/util/assert';
@@ -270,6 +270,10 @@ export class TestPlatform implements Platform {
   randomBytes(nBytes: number): Uint8Array {
     return this.basePlatform.randomBytes(nBytes);
   }
+
+  toByteStreamReader(source: unknown): ByteStreamReader {
+    return this.basePlatform.toByteStreamReader(source);
+  }
 }
 
 /** Returns true if we are running under Node. */