Implement bundle loading.

packages/firestore-types/index.d.ts

@@ -93,9 +93,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?: () => void
+  ): Promise<void>;
+
+  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/bundle.ts

@@ -0,0 +1,108 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as firestore from '@firebase/firestore-types';
+import { Deferred } from '../util/promise';
+
+export class LoadBundleTask implements firestore.LoadBundleTask {
+  private _progressResolver = new Deferred<void>();
+  private _userProgressHandler?: (
+    progress: firestore.LoadBundleTaskProgress
+  ) => unknown;
+  private _userProgressErrorHandler?: (err: Error) => unknown;
+  private _userProgressCompleteHandler?: () => void;
+
+  private _promiseResolver = new Deferred<firestore.LoadBundleTaskProgress>();
+
+  private _lastProgress: firestore.LoadBundleTaskProgress = {
+    taskState: 'Running',
+    totalBytes: 0,
+    totalDocuments: 0,
+    bytesLoaded: 0,
+    documentsLoaded: 0
+  };
+
+  onProgress(
+    next?: (progress: firestore.LoadBundleTaskProgress) => unknown,
+    error?: (err: Error) => unknown,
+    complete?: () => void
+  ): Promise<void> {
+    this._userProgressHandler = next;
+    this._userProgressErrorHandler = error;
+    this._userProgressCompleteHandler = complete;
+    return this._progressResolver.promise;
+  }
+
+  catch(onRejected: (a: Error) => unknown): Promise<unknown> {
+    return this._promiseResolver.promise.catch(onRejected);
+  }
+
+  then(
+    onFulfilled?: (a: firestore.LoadBundleTaskProgress) => unknown,
+    onRejected?: (a: Error) => unknown
+  ): Promise<unknown> {
+    return this._promiseResolver.promise.then(onFulfilled, onRejected);
+  }
+
+  /**
+   * Notifies the completion of loading a bundle, with a provided
+   * `LoadBundleTaskProgress` object.
+   */
+  _completeWith(progress: firestore.LoadBundleTaskProgress): void {
+    this._updateProgress(progress);
+    if (this._userProgressCompleteHandler) {
+      this._userProgressCompleteHandler();
+    }
+    this._progressResolver.resolve();
+
+    this._promiseResolver.resolve(progress);
+  }
+
+  /**
+   * Notifies a failure of loading a bundle, with a provided `Error`
+   * as the reason.
+   */
+  _failedWith(error: Error): void {
+    this._lastProgress.taskState = 'Error';
+
+    if (this._userProgressHandler) {
+      this._userProgressHandler(this._lastProgress);
+    }
+
+    if (this._userProgressErrorHandler) {
+      this._userProgressErrorHandler(error);
+    }
+    this._progressResolver.reject(error);
+
+    this._promiseResolver.reject(error);
+  }
+
+  /**
+   * Notifies a progress update of loading a bundle.
+   * @param progress The new progress.
+   */
+  _updateProgress(progress: firestore.LoadBundleTaskProgress): void {
+    if (this._lastProgress.taskState === 'Error') {
+      return;
+    }
+
+    this._lastProgress = progress;
+    if (this._userProgressHandler) {
+      this._userProgressHandler(progress);
+    }
+  }
+}

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,14 @@ 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 { BundleMetadata } from '../protos/firestore_bundle_proto';
 
 /**
  * Represents a Firestore bundle saved by the SDK in its local storage.
@@ -98,3 +107,147 @@ export class BundleConverter {
     return fromVersion(time);
   }
 }
+
+/**
+ * Returns a `LoadBundleTaskProgress` representing the initial progress of
+ * loading a bundle.
+ */
+export function initialProgress(
+  metadata: BundleMetadata
+): firestore.LoadBundleTaskProgress {
+  return {
+    taskState: 'Running',
+    documentsLoaded: 0,
+    bytesLoaded: 0,
+    totalDocuments: metadata.totalDocuments!,
+    totalBytes: metadata.totalBytes!
+  };
+}
+
+/**
+ * Returns a `LoadBundleTaskProgress` representing the progress if the bundle
+ * is already loaded, and we are skipping current loading.
+ */
+export function skipLoadingProgress(
+  metadata: BundleMetadata
+): firestore.LoadBundleTaskProgress {
+  return {
+    taskState: 'Success',
+    documentsLoaded: metadata.totalDocuments!,
+    bytesLoaded: metadata.totalBytes!,
+    totalDocuments: metadata.totalDocuments!,
+    totalBytes: metadata.totalBytes!
+  };
+}
+
+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.
+   *
+   * Applies to both `documentsBuffered` and `bytesBuffered`, triggers storage
+   * update and reports progress when either of them cross the threshold.
+   */
+  private thresholdMultiplier = 0.01;
+  /** Batched queries to be saved into storage */
+  private queries: bundleProto.NamedQuery[] = [];
+  /** Batched documents to be saved into storage */
+  private documents: BundledDocuments = [];
+  /**
+   * 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(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
+  ): firestore.LoadBundleTaskProgress | null {
+    debugAssert(!element.isBundleMetadata(), 'Unexpected bundle metadata.');
+
+    this.progress.bytesLoaded += 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.progress.documentsLoaded += 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.progress.documentsLoaded += 1;
+      this.unpairedDocumentMetadata = null;
+    }
+
+    // Loading a document metadata will not update progress.
+    if (this.unpairedDocumentMetadata) {
+      return null;
+    }
+
+    return { ...this.progress };
+  }
+
+  /**
+   * Update the progress to 'Success' and return the updated progress.
+   */
+  async complete(): Promise<LoadResult> {
+    debugAssert(
+      !this.unpairedDocumentMetadata,
+      'Unexpected document when no pairing metadata is found'
+    );
+
+    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.taskState = 'Success';
+    return new LoadResult({ ...this.progress }, changedDocs);
+  }
+}

packages/firestore/src/core/firestore_client.ts

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+import * as firestore 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 { LoadBundleTask } from '../api/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
+  ): firestore.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 LoadBundleTask();
+    this.asyncQueue.enqueueAndForget(() => {
+      return loadBundle(this.syncEngine, reader, task);
+    });
+
+    return task;
+  }
 }