Change ProtoByteString to use its own class rather than String | Uint8Array (#2639)

Author: Brian Chen

packages/firestore/src/api/blob.ts

@@ -24,6 +24,10 @@ import {
   validateExactNumberOfArgs
 } from '../util/input_validation';
 import { primitiveComparator } from '../util/misc';
+import {
+  binaryStringFromUint8Array,
+  uint8ArrayFromBinaryString
+} from '../util/proto_byte_string';
 
 /** Helper function to assert Uint8Array is available at runtime. */
 function assertUint8ArrayAvailable(): void {
@@ -85,10 +89,7 @@ export class Blob {
     if (!(array instanceof Uint8Array)) {
       throw invalidClassError('Blob.fromUint8Array', 'Uint8Array', 1, array);
     }
-    let binaryString = '';
-    for (let i = 0; i < array.length; ++i) {
-      binaryString += String.fromCharCode(array[i]);
-    }
+    const binaryString = binaryStringFromUint8Array(array);
     return new Blob(binaryString);
   }
 
@@ -101,10 +102,7 @@ export class Blob {
   toUint8Array(): Uint8Array {
     validateExactNumberOfArgs('Blob.toUint8Array', arguments, 0);
     assertUint8ArrayAvailable();
-    const buffer = new Uint8Array(this._binaryString.length);
-    for (let i = 0; i < this._binaryString.length; i++) {
-      buffer[i] = this._binaryString.charCodeAt(i);
-    }
+    const buffer = uint8ArrayFromBinaryString(this._binaryString);
     return buffer;
   }
 

packages/firestore/src/core/types.ts

@@ -29,10 +29,6 @@ export type TargetId = number;
 
 export type ListenSequenceNumber = number;
 
-// TODO(b/35918695): In GRPC / node, tokens are Uint8Array. In WebChannel,
-// they're strings. We should probably (de-)serialize to a common internal type.
-export type ProtoByteString = Uint8Array | string;
-
 /** The different states of a mutation batch. */
 export type MutationBatchState = 'pending' | 'acknowledged' | 'rejected';
 

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -18,14 +18,15 @@
 import { Timestamp } from '../api/timestamp';
 import { User } from '../auth/user';
 import { Query } from '../core/query';
-import { BatchId, ProtoByteString } from '../core/types';
+import { BatchId } from '../core/types';
 import { DocumentKeySet } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
 import { ResourcePath } from '../model/path';
 import { assert, fail } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { ByteString } from '../util/proto_byte_string';
 import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
@@ -47,7 +48,7 @@ import { LocalSerializer } from './local_serializer';
 import { MutationQueue } from './mutation_queue';
 import { PersistenceTransaction, ReferenceDelegate } from './persistence';
 import { PersistencePromise } from './persistence_promise';
-import { SimpleDb, SimpleDbStore, SimpleDbTransaction } from './simple_db';
+import { SimpleDbStore, SimpleDbTransaction } from './simple_db';
 
 /** A mutation queue for a specific user, backed by IndexedDB. */
 export class IndexedDbMutationQueue implements MutationQueue {
@@ -121,29 +122,33 @@ export class IndexedDbMutationQueue implements MutationQueue {
   acknowledgeBatch(
     transaction: PersistenceTransaction,
     batch: MutationBatch,
-    streamToken: ProtoByteString
+    streamToken: ByteString
   ): PersistencePromise<void> {
     return this.getMutationQueueMetadata(transaction).next(metadata => {
-      metadata.lastStreamToken = convertStreamToken(streamToken);
+      // We can't store the resumeToken as a ByteString in IndexedDB, so we
+      // convert it to a Base64 string for storage.
+      metadata.lastStreamToken = streamToken.toBase64();
 
       return mutationQueuesStore(transaction).put(metadata);
     });
   }
 
   getLastStreamToken(
     transaction: PersistenceTransaction
-  ): PersistencePromise<ProtoByteString> {
-    return this.getMutationQueueMetadata(transaction).next<ProtoByteString>(
-      metadata => metadata.lastStreamToken
+  ): PersistencePromise<ByteString> {
+    return this.getMutationQueueMetadata(transaction).next<ByteString>(
+      metadata => ByteString.fromBase64String(metadata.lastStreamToken)
     );
   }
 
   setLastStreamToken(
     transaction: PersistenceTransaction,
-    streamToken: ProtoByteString
+    streamToken: ByteString
   ): PersistencePromise<void> {
     return this.getMutationQueueMetadata(transaction).next(metadata => {
-      metadata.lastStreamToken = convertStreamToken(streamToken);
+      // We can't store the resumeToken as a ByteString in IndexedDB, so we
+      // convert it to a Base64 string for storage.
+      metadata.lastStreamToken = streamToken.toBase64();
       return mutationQueuesStore(transaction).put(metadata);
     });
   }
@@ -671,19 +676,6 @@ export function removeMutationBatch(
   return PersistencePromise.waitFor(promises).next(() => removedDocuments);
 }
 
-function convertStreamToken(token: ProtoByteString): string {
-  if (token instanceof Uint8Array) {
-    // TODO(b/78771403): Convert tokens to strings during deserialization
-    assert(
-      SimpleDb.isMockPersistence(),
-      'Persisting non-string stream tokens is only supported with mock persistence.'
-    );
-    return token.toString();
-  } else {
-    return token;
-  }
-}
-
 /**
  * Helper to get a typed SimpleDbStore for the mutations object store.
  */

packages/firestore/src/local/local_serializer.ts

@@ -28,6 +28,7 @@ import { MutationBatch } from '../model/mutation_batch';
 import * as api from '../protos/firestore_proto_api';
 import { JsonProtoSerializer } from '../remote/serializer';
 import { assert, fail } from '../util/assert';
+import { ByteString } from '../util/proto_byte_string';
 
 import { documentKeySet, DocumentKeySet } from '../model/collections';
 import { Target } from '../core/target';
@@ -42,7 +43,6 @@ import {
   DbTimestampKey,
   DbUnknownDocument
 } from './indexeddb_schema';
-import { SimpleDb } from './simple_db';
 import { TargetData, TargetPurpose } from './target_data';
 
 /** Serializer for values stored in the LocalStore. */
@@ -207,8 +207,6 @@ export class LocalSerializer {
       dbTarget.lastLimboFreeSnapshotVersion !== undefined
         ? this.fromDbTimestamp(dbTarget.lastLimboFreeSnapshotVersion)
         : SnapshotVersion.MIN;
-    // TODO(b/140573486): Convert to platform representation
-    const resumeToken = dbTarget.resumeToken;
 
     let target: Target;
     if (isDocumentQuery(dbTarget.query)) {
@@ -223,7 +221,7 @@ export class LocalSerializer {
       dbTarget.lastListenSequenceNumber,
       version,
       lastLimboFreeSnapshotVersion,
-      resumeToken
+      ByteString.fromBase64String(dbTarget.resumeToken)
     );
   }
 
@@ -247,18 +245,9 @@ export class LocalSerializer {
       queryProto = this.remoteSerializer.toQueryTarget(targetData.target);
     }
 
-    let resumeToken: string;
-
-    if (targetData.resumeToken instanceof Uint8Array) {
-      // TODO(b/78771403): Convert tokens to strings during deserialization
-      assert(
-        SimpleDb.isMockPersistence(),
-        'Persisting non-string stream tokens is only supported with mock persistence .'
-      );
-      resumeToken = targetData.resumeToken.toString();
-    } else {
-      resumeToken = targetData.resumeToken;
-    }
+    // We can't store the resumeToken as a ByteString in IndexedDb, so we
+    // convert it to a base64 string for storage.
+    const resumeToken = targetData.resumeToken.toBase64();
 
     // lastListenSequenceNumber is always 0 until we do real GC.
     return new DbTarget(

packages/firestore/src/local/local_store.ts

@@ -20,7 +20,7 @@ import { User } from '../auth/user';
 import { Query } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
 import { Target } from '../core/target';
-import { BatchId, ProtoByteString, TargetId } from '../core/types';
+import { BatchId, TargetId } from '../core/types';
 import {
   DocumentKeySet,
   documentKeySet,
@@ -62,6 +62,7 @@ import { RemoteDocumentCache } from './remote_document_cache';
 import { RemoteDocumentChangeBuffer } from './remote_document_change_buffer';
 import { ClientId } from './shared_client_state';
 import { TargetData, TargetPurpose } from './target_data';
+import { ByteString } from '../util/proto_byte_string';
 
 const LOG_TAG = 'LocalStore';
 
@@ -464,7 +465,7 @@ export class LocalStore {
   }
 
   /** Returns the last recorded stream token for the current user. */
-  getLastStreamToken(): Promise<ProtoByteString> {
+  getLastStreamToken(): Promise<ByteString> {
     return this.persistence.runTransaction(
       'Get last stream token',
       'readonly-idempotent',
@@ -479,7 +480,7 @@ export class LocalStore {
    * mutation batch. This is usually only useful after a stream handshake or in
    * response to an error that requires clearing the stream token.
    */
-  setLastStreamToken(streamToken: ProtoByteString): Promise<void> {
+  setLastStreamToken(streamToken: ByteString): Promise<void> {
     return this.persistence.runTransaction(
       'Set last stream token',
       'readwrite-primary-idempotent',
@@ -551,7 +552,7 @@ export class LocalStore {
 
               const resumeToken = change.resumeToken;
               // Update the resume token if the change includes one.
-              if (resumeToken.length > 0) {
+              if (resumeToken.approximateByteSize() > 0) {
                 const newTargetData = oldTargetData
                   .withResumeToken(resumeToken, remoteVersion)
                   .withSequenceNumber(txn.currentSequenceNumber);
@@ -696,12 +697,12 @@ export class LocalStore {
     change: TargetChange
   ): boolean {
     assert(
-      newTargetData.resumeToken.length > 0,
+      newTargetData.resumeToken.approximateByteSize() > 0,
       'Attempted to persist target data with no resume token'
     );
 
     // Always persist target data if we don't already have a resume token.
-    if (oldTargetData.resumeToken.length === 0) {
+    if (oldTargetData.resumeToken.approximateByteSize() === 0) {
       return true;
     }
 

packages/firestore/src/local/memory_mutation_queue.ts

@@ -17,14 +17,14 @@
 
 import { Timestamp } from '../api/timestamp';
 import { Query } from '../core/query';
-import { BatchId, ProtoByteString } from '../core/types';
+import { BatchId } from '../core/types';
 import { DocumentKeySet } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { MutationBatch, BATCHID_UNKNOWN } from '../model/mutation_batch';
-import { emptyByteString } from '../platform/platform';
 import { assert } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { ByteString } from '../util/proto_byte_string';
 import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
@@ -48,7 +48,7 @@ export class MemoryMutationQueue implements MutationQueue {
    * responses the client has processed. Stream tokens are opaque checkpoint
    * markers whose only real value is their inclusion in the next request.
    */
-  private lastStreamToken: ProtoByteString = emptyByteString();
+  private lastStreamToken: ByteString = ByteString.EMPTY_BYTE_STRING;
 
   /** An ordered mapping between documents and the mutations batch IDs. */
   private batchesByDocumentKey = new SortedSet(DocReference.compareByKey);
@@ -65,7 +65,7 @@ export class MemoryMutationQueue implements MutationQueue {
   acknowledgeBatch(
     transaction: PersistenceTransaction,
     batch: MutationBatch,
-    streamToken: ProtoByteString
+    streamToken: ByteString
   ): PersistencePromise<void> {
     const batchId = batch.batchId;
     const batchIndex = this.indexOfExistingBatchId(batchId, 'acknowledged');
@@ -90,13 +90,13 @@ export class MemoryMutationQueue implements MutationQueue {
 
   getLastStreamToken(
     transaction: PersistenceTransaction
-  ): PersistencePromise<ProtoByteString> {
+  ): PersistencePromise<ByteString> {
     return PersistencePromise.resolve(this.lastStreamToken);
   }
 
   setLastStreamToken(
     transaction: PersistenceTransaction,
-    streamToken: ProtoByteString
+    streamToken: ByteString
   ): PersistencePromise<void> {
     this.lastStreamToken = streamToken;
     return PersistencePromise.resolve();

packages/firestore/src/local/mutation_queue.ts

@@ -17,11 +17,12 @@
 
 import { Timestamp } from '../api/timestamp';
 import { Query } from '../core/query';
-import { BatchId, ProtoByteString } from '../core/types';
+import { BatchId } from '../core/types';
 import { DocumentKeySet } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { MutationBatch } from '../model/mutation_batch';
+import { ByteString } from '../util/proto_byte_string';
 import { SortedMap } from '../util/sorted_map';
 
 import { PersistenceTransaction } from './persistence';
@@ -38,18 +39,18 @@ export interface MutationQueue {
   acknowledgeBatch(
     transaction: PersistenceTransaction,
     batch: MutationBatch,
-    streamToken: ProtoByteString
+    streamToken: ByteString
   ): PersistencePromise<void>;
 
   /** Returns the current stream token for this mutation queue. */
   getLastStreamToken(
     transaction: PersistenceTransaction
-  ): PersistencePromise<ProtoByteString>;
+  ): PersistencePromise<ByteString>;
 
   /** Sets the stream token for this mutation queue. */
   setLastStreamToken(
     transaction: PersistenceTransaction,
-    streamToken: ProtoByteString
+    streamToken: ByteString
   ): PersistencePromise<void>;
 
   /**

packages/firestore/src/local/target_data.ts

@@ -17,8 +17,8 @@
 
 import { SnapshotVersion } from '../core/snapshot_version';
 import { Target } from '../core/target';
-import { ListenSequenceNumber, ProtoByteString, TargetId } from '../core/types';
-import { emptyByteString } from '../platform/platform';
+import { ListenSequenceNumber, TargetId } from '../core/types';
+import { ByteString } from '../util/proto_byte_string';
 
 /** An enumeration of the different purposes we have for targets. */
 export enum TargetPurpose {
@@ -66,7 +66,7 @@ export class TargetData {
      * matches the target. The resume token essentially identifies a point in
      * time from which the server should resume sending results.
      */
-    readonly resumeToken: ProtoByteString = emptyByteString()
+    readonly resumeToken: ByteString = ByteString.EMPTY_BYTE_STRING
   ) {}
 
   /** Creates a new target data instance with an updated sequence number. */
@@ -87,7 +87,7 @@ export class TargetData {
    * snapshot version.
    */
   withResumeToken(
-    resumeToken: ProtoByteString,
+    resumeToken: ByteString,
     snapshotVersion: SnapshotVersion
   ): TargetData {
     return new TargetData(
@@ -128,7 +128,7 @@ export class TargetData {
       this.lastLimboFreeSnapshotVersion.isEqual(
         other.lastLimboFreeSnapshotVersion
       ) &&
-      this.resumeToken === other.resumeToken &&
+      this.resumeToken.isEqual(other.resumeToken) &&
       this.target.isEqual(other.target)
     );
   }