Implement bundle loading.

packages/firestore-types/index.d.ts

@@ -92,9 +92,38 @@ export class FirebaseFirestore {
 
   terminate(): Promise<void>;
 
+  loadBundle(
+    bundleData: ArrayBuffer | ReadableStream<ArrayBuffer> | string
+  ): LoadBundleTask;
+
   INTERNAL: { delete: () => Promise<void> };
 }
 
+export interface LoadBundleTask {
+  onProgress(
+    next?: (progress: LoadBundleTaskProgress) => any,
+    error?: (error: Error) => any,
+    complete?: (progress?: LoadBundleTaskProgress) => any
+  ): Promise<any>;
+
+  then(
+    onFulfilled?: (a: LoadBundleTaskProgress) => any,
+    onRejected?: (a: Error) => any
+  ): Promise<any>;
+
+  catch(onRejected: (a: Error) => any): Promise<any>;
+}
+
+export interface LoadBundleTaskProgress {
+  documentsLoaded: number;
+  totalDocuments: number;
+  bytesLoaded: number;
+  totalBytes: number;
+  taskState: TaskState;
+}
+
+export type TaskState = 'Error' | 'Running' | 'Success';
+
 export class GeoPoint {
   constructor(latitude: number, longitude: number);
 

packages/firestore/src/api/database.ts

@@ -493,6 +493,13 @@ export class Firestore implements firestore.FirebaseFirestore, FirebaseService {
     };
   }
 
+  loadBundle(
+    bundleData: ArrayBuffer | ReadableStream<Uint8Array> | string
+  ): firestore.LoadBundleTask {
+    this.ensureClientConfigured();
+    return this._firestoreClient!.loadBundle(bundleData);
+  }
+
   ensureClientConfigured(): FirestoreClient {
     if (!this._firestoreClient) {
       // Kick off starting the client but don't actually wait for it.

packages/firestore/src/core/bundle.ts

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+import * as firestore from '@firebase/firestore-types';
 import { Query } from './query';
 import { SnapshotVersion } from './snapshot_version';
 import {
@@ -28,6 +29,15 @@ import * as api from '../protos/firestore_proto_api';
 import { DocumentKey } from '../model/document_key';
 import { MaybeDocument, NoDocument } from '../model/document';
 import { debugAssert } from '../util/assert';
+import {
+  applyBundleDocuments,
+  LocalStore,
+  saveNamedQuery
+} from '../local/local_store';
+import { SizedBundleElement } from '../util/bundle_reader';
+import { MaybeDocumentMap } from '../model/collections';
+import { Deferred } from '../util/promise';
+import { BundleMetadata } from '../protos/firestore_bundle_proto';
 
 /**
  * Represents a Firestore bundle saved by the SDK in its local storage.
@@ -98,3 +108,247 @@ export class BundleConverter {
     return fromVersion(time);
   }
 }
+
+/**
+ * Returns a `LoadBundleTaskProgress` representing the first progress of
+ * loading a bundle.
+ */
+export function initialProgress(
+  state: firestore.TaskState,
+  metadata: BundleMetadata
+): firestore.LoadBundleTaskProgress {
+  return {
+    taskState: state,
+    documentsLoaded: state === 'Success' ? metadata.totalDocuments! : 0,
+    bytesLoaded: state === 'Success' ? metadata.totalBytes! : 0,
+    totalDocuments: metadata.totalDocuments!,
+    totalBytes: metadata.totalBytes!
+  };
+}
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+export class LoadBundleTaskImpl implements firestore.LoadBundleTask {
+  private progressResolver = new Deferred<any>();
+  private progressNext?: (progress: firestore.LoadBundleTaskProgress) => any;
+  private progressError?: (err: Error) => any;
+  private progressComplete?: (
+    progress?: firestore.LoadBundleTaskProgress
+  ) => any;
+
+  private promiseResolver = new Deferred<any>();
+  private promiseFulfilled?: (
+    progress: firestore.LoadBundleTaskProgress
+  ) => any;
+  private promiseRejected?: (err: Error) => any;
+
+  private lastProgress: firestore.LoadBundleTaskProgress = {
+    taskState: 'Running',
+    totalBytes: 0,
+    totalDocuments: 0,
+    bytesLoaded: 0,
+    documentsLoaded: 0
+  };
+
+  onProgress(
+    next?: (progress: firestore.LoadBundleTaskProgress) => any,
+    error?: (err: Error) => any,
+    complete?: (progress?: firestore.LoadBundleTaskProgress) => void
+  ): Promise<any> {
+    this.progressNext = next;
+    this.progressError = error;
+    this.progressComplete = complete;
+    return this.progressResolver.promise;
+  }
+
+  catch(onRejected: (a: Error) => any): Promise<any> {
+    this.promiseRejected = onRejected;
+    return this.promiseResolver.promise;
+  }
+
+  then(
+    onFulfilled?: (a: firestore.LoadBundleTaskProgress) => any,
+    onRejected?: (a: Error) => any
+  ): Promise<any> {
+    this.promiseFulfilled = onFulfilled;
+    this.promiseRejected = onRejected;
+    return this.promiseResolver.promise;
+  }
+  /* eslint-enable @typescript-eslint/no-explicit-any */
+
+  /**
+   * Notifies the completion of loading a bundle, with a provided
+   * `LoadBundleTaskProgress` object.
+   */
+  completeWith(progress: firestore.LoadBundleTaskProgress): void {
+    let result;
+    if (this.progressComplete) {
+      result = this.progressComplete(progress);
+    }
+    this.progressResolver.resolve(result);
+
+    result = undefined;
+    if (this.promiseFulfilled) {
+      result = this.promiseFulfilled(progress);
+    }
+    this.promiseResolver.resolve(result);
+  }
+
+  /**
+   * Notifies a failure of loading a bundle, with a provided `Error`
+   * as the reason.
+   */
+  failedWith(error: Error): void {
+    if (this.progressNext) {
+      this.lastProgress.taskState = 'Error';
+      this.progressNext(this.lastProgress);
+    }
+
+    let result;
+    if (this.progressError) {
+      result = this.progressError(error);
+    }
+    this.progressResolver.reject(result);
+
+    result = undefined;
+    if (this.promiseRejected) {
+      this.promiseRejected(error);
+    }
+    this.promiseResolver.reject(result);
+  }
+
+  /**
+   * Notifies a progress update of loading a bundle.
+   * @param progress The new progress.
+   */
+  updateProgress(progress: firestore.LoadBundleTaskProgress): void {
+    this.lastProgress = progress;
+    if (this.progressNext) {
+      this.progressNext(progress);
+    }
+  }
+}
+
+export class LoadResult {
+  constructor(
+    readonly progress: firestore.LoadBundleTaskProgress,
+    readonly changedDocs?: MaybeDocumentMap
+  ) {}
+}
+
+/**
+ * A class to process the elements from a bundle, load them into local
+ * storage and provide progress update while loading.
+ */
+export class BundleLoader {
+  /** The current progress of loading */
+  private progress: firestore.LoadBundleTaskProgress;
+  /**
+   * The threshold multiplier used to determine whether enough elements are
+   * batched to be loaded, and a progress update is needed.
+   */
+  private step = 0.01;
+  /** Batched queries to be saved into storage */
+  private queries: bundleProto.NamedQuery[] = [];
+  /** Batched documents to be saved into storage */
+  private documents: BundledDocuments = [];
+  /** How many bytes in the bundle are being batched. */
+  private bytesIncrement = 0;
+  /** How many documents in the bundle are being batched. */
+  private documentsIncrement = 0;
+  /**
+   * A BundleDocumentMetadata is added to the loader, it is saved here while
+   * we wait for the actual document.
+   */
+  private unpairedDocumentMetadata: bundleProto.BundledDocumentMetadata | null = null;
+
+  constructor(
+    private metadata: bundleProto.BundleMetadata,
+    private localStore: LocalStore
+  ) {
+    this.progress = initialProgress('Running', metadata);
+  }
+
+  /**
+   * Adds an element from the bundle to the loader.
+   *
+   * If adding this element leads to actually saving the batched elements into
+   * storage, the returned promise will resolve to a `LoadResult`, otherwise
+   * it will resolve to null.
+   */
+  addSizedElement(element: SizedBundleElement): Promise<LoadResult | null> {
+    debugAssert(!element.isBundleMetadata(), 'Unexpected bundle metadata.');
+
+    this.bytesIncrement += element.byteLength;
+    if (element.payload.namedQuery) {
+      this.queries.push(element.payload.namedQuery);
+    }
+
+    if (element.payload.documentMetadata) {
+      if (element.payload.documentMetadata.exists) {
+        this.unpairedDocumentMetadata = element.payload.documentMetadata;
+      } else {
+        this.documents.push({
+          metadata: element.payload.documentMetadata,
+          document: undefined
+        });
+        this.documentsIncrement += 1;
+      }
+    }
+
+    if (element.payload.document) {
+      debugAssert(
+        !!this.unpairedDocumentMetadata,
+        'Unexpected document when no pairing metadata is found'
+      );
+      this.documents.push({
+        metadata: this.unpairedDocumentMetadata!,
+        document: element.payload.document
+      });
+      this.documentsIncrement += 1;
+      this.unpairedDocumentMetadata = null;
+    }
+
+    return this.saveAndReportProgress();
+  }
+
+  private async saveAndReportProgress(): Promise<LoadResult | null> {
+    if (
+      this.unpairedDocumentMetadata ||
+      (this.documentsIncrement < this.progress.totalDocuments * this.step &&
+        this.bytesIncrement < this.progress.totalBytes * this.step)
+    ) {
+      return null;
+    }
+
+    for (const q of this.queries) {
+      await saveNamedQuery(this.localStore, q);
+    }
+
+    let changedDocs;
+    if (this.documents.length > 0) {
+      changedDocs = await applyBundleDocuments(this.localStore, this.documents);
+    }
+
+    this.progress.bytesLoaded += this.bytesIncrement;
+    this.progress.documentsLoaded += this.documentsIncrement;
+    this.bytesIncrement = 0;
+    this.documentsIncrement = 0;
+    this.queries = [];
+    this.documents = [];
+
+    return new LoadResult({ ...this.progress }, changedDocs);
+  }
+
+  /**
+   * Update the progress to 'Success' and return the updated progress.
+   */
+  complete(): firestore.LoadBundleTaskProgress {
+    debugAssert(
+      this.queries.length === 0 && this.documents.length === 0,
+      'There are more items needs to be saved but complete() is called.'
+    );
+    this.progress.taskState = 'Success';
+
+    return this.progress;
+  }
+}

packages/firestore/src/core/firestore_client.ts

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+import { LoadBundleTask } from '@firebase/firestore-types';
 import { CredentialsProvider } from '../api/credentials';
 import { User } from '../auth/user';
 import { LocalStore } from '../local/local_store';
@@ -34,7 +35,7 @@ import {
   Observer,
   QueryListener
 } from './event_manager';
-import { SyncEngine } from './sync_engine';
+import { SyncEngine, loadBundle } from './sync_engine';
 import { View } from './view';
 
 import { SharedClientState } from '../local/shared_client_state';
@@ -47,8 +48,11 @@ import {
   ComponentProvider,
   MemoryComponentProvider
 } from './component_provider';
+import { BundleReader } from '../util/bundle_reader';
+import { LoadBundleTaskImpl } from './bundle';
 import { newConnection } from '../platform/connection';
 import { newSerializer } from '../platform/serializer';
+import { toByteStreamReader } from '../platform/byte_stream_reader';
 
 const LOG_TAG = 'FirestoreClient';
 const MAX_CONCURRENT_LIMBO_RESOLUTIONS = 100;
@@ -512,4 +516,24 @@ export class FirestoreClient {
     });
     return deferred.promise;
   }
+
+  loadBundle(
+    data: ReadableStream<Uint8Array> | ArrayBuffer | string
+  ): LoadBundleTask {
+    this.verifyNotTerminated();
+
+    let content: ReadableStream<Uint8Array> | ArrayBuffer;
+    if (typeof data === 'string') {
+      content = new TextEncoder().encode(data);
+    } else {
+      content = data;
+    }
+    const reader = new BundleReader(toByteStreamReader(content));
+    const task = new LoadBundleTaskImpl();
+    this.asyncQueue.enqueueAndForget(() => {
+      return loadBundle(this.syncEngine, reader, task);
+    });
+
+    return task;
+  }
 }