Implement bundle loading. (#3201)

Author: wu-hui

packages/firestore-types/index.d.ts

@@ -93,9 +93,40 @@ 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
+  ): void;
+
+  then<T, R>(
+    onFulfilled?: (a: LoadBundleTaskProgress) => T | PromiseLike<T>,
+    onRejected?: (a: Error) => R | PromiseLike<R>
+  ): Promise<T | R>;
+
+  catch<R>(
+    onRejected: (a: Error) => R | PromiseLike<R>
+  ): Promise<R | LoadBundleTaskProgress>;
+}
+
+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,117 @@
+/**
+ * @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';
+import { PartialObserver } from './observer';
+import { debugAssert } from '../util/assert';
+
+export class LoadBundleTask
+  implements
+    firestore.LoadBundleTask,
+    PromiseLike<firestore.LoadBundleTaskProgress> {
+  private _progressObserver: PartialObserver<
+    firestore.LoadBundleTaskProgress
+  > = {};
+  private _taskCompletionResolver = 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
+  ): void {
+    this._progressObserver = {
+      next,
+      error,
+      complete
+    };
+  }
+
+  catch<R>(
+    onRejected: (a: Error) => R | PromiseLike<R>
+  ): Promise<R | firestore.LoadBundleTaskProgress> {
+    return this._taskCompletionResolver.promise.catch(onRejected);
+  }
+
+  then<T, R>(
+    onFulfilled?: (a: firestore.LoadBundleTaskProgress) => T | PromiseLike<T>,
+    onRejected?: (a: Error) => R | PromiseLike<R>
+  ): Promise<T | R> {
+    return this._taskCompletionResolver.promise.then(onFulfilled, onRejected);
+  }
+
+  /**
+   * Notifies all observers that bundle loading has completed, with a provided
+   * `LoadBundleTaskProgress` object.
+   */
+  _completeWith(progress: firestore.LoadBundleTaskProgress): void {
+    debugAssert(
+      progress.taskState === 'Success',
+      'Task is not completed with Success.'
+    );
+    this._updateProgress(progress);
+    if (this._progressObserver.complete) {
+      this._progressObserver.complete();
+    }
+
+    this._taskCompletionResolver.resolve(progress);
+  }
+
+  /**
+   * Notifies all observers that bundle loading has failed, with a provided
+   * `Error` as the reason.
+   */
+  _failWith(error: Error): void {
+    this._lastProgress.taskState = 'Error';
+
+    if (this._progressObserver.next) {
+      this._progressObserver.next(this._lastProgress);
+    }
+
+    if (this._progressObserver.error) {
+      this._progressObserver.error(error);
+    }
+
+    this._taskCompletionResolver.reject(error);
+  }
+
+  /**
+   * Notifies a progress update of loading a bundle.
+   * @param progress The new progress.
+   */
+  _updateProgress(progress: firestore.LoadBundleTaskProgress): void {
+    debugAssert(
+      this._lastProgress.taskState === 'Running',
+      'Cannot update progress on a completed or failed task'
+    );
+
+    this._lastProgress = progress;
+    if (this._progressObserver.next) {
+      this._progressObserver.next(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.
@@ -58,7 +67,7 @@ export interface NamedQuery {
  */
 interface BundledDocument {
   metadata: bundleProto.BundledDocumentMetadata;
-  document: api.Document | undefined;
+  document?: api.Document;
 }
 
 /**
@@ -98,3 +107,127 @@ export class BundleConverter {
     return fromVersion(time);
   }
 }
+
+/**
+ * Returns a `LoadBundleTaskProgress` representing the initial progress of
+ * loading a bundle.
+ */
+export function bundleInitialProgress(
+  metadata: BundleMetadata
+): firestore.LoadBundleTaskProgress {
+  return {
+    taskState: 'Running',
+    documentsLoaded: 0,
+    bytesLoaded: 0,
+    totalDocuments: metadata.totalDocuments!,
+    totalBytes: metadata.totalBytes!
+  };
+}
+
+/**
+ * Returns a `LoadBundleTaskProgress` representing the progress that the loading
+ * has succeeded.
+ */
+export function bundleSuccessProgress(
+  metadata: BundleMetadata
+): firestore.LoadBundleTaskProgress {
+  return {
+    taskState: 'Success',
+    documentsLoaded: metadata.totalDocuments!,
+    bytesLoaded: metadata.totalBytes!,
+    totalDocuments: metadata.totalDocuments!,
+    totalBytes: metadata.totalBytes!
+  };
+}
+
+export class BundleLoadResult {
+  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;
+  /** Batched queries to be saved into storage */
+  private queries: bundleProto.NamedQuery[] = [];
+  /** Batched documents to be saved into storage */
+  private documents: BundledDocuments = [];
+
+  constructor(
+    private metadata: bundleProto.BundleMetadata,
+    private localStore: LocalStore
+  ) {
+    this.progress = bundleInitialProgress(metadata);
+  }
+
+  /**
+   * Adds an element from the bundle to the loader.
+   *
+   * Returns a new progress if adding the element leads to a new progress,
+   * otherwise returns null.
+   */
+  addSizedElement(
+    element: SizedBundleElement
+  ): firestore.LoadBundleTaskProgress | null {
+    debugAssert(!element.isBundleMetadata(), 'Unexpected bundle metadata.');
+
+    this.progress.bytesLoaded += element.byteLength;
+
+    let documentsLoaded = this.progress.documentsLoaded;
+
+    if (element.payload.namedQuery) {
+      this.queries.push(element.payload.namedQuery);
+    } else if (element.payload.documentMetadata) {
+      this.documents.push({ metadata: element.payload.documentMetadata });
+      if (!element.payload.documentMetadata.exists) {
+        ++documentsLoaded;
+      }
+    } else if (element.payload.document) {
+      debugAssert(
+        this.documents.length > 0 &&
+          this.documents[this.documents.length - 1].metadata.name ===
+            element.payload.document.name,
+        'The document being added does not match the stored metadata.'
+      );
+      this.documents[this.documents.length - 1].document =
+        element.payload.document;
+      ++documentsLoaded;
+    }
+
+    if (documentsLoaded !== this.progress.documentsLoaded) {
+      this.progress.documentsLoaded = documentsLoaded;
+      return { ...this.progress };
+    }
+
+    return null;
+  }
+
+  /**
+   * Update the progress to 'Success' and return the updated progress.
+   */
+  async complete(): Promise<BundleLoadResult> {
+    debugAssert(
+      this.documents[this.documents.length - 1]?.metadata.exists !== true ||
+        !!this.documents[this.documents.length - 1].document,
+      'Bundled documents ends with a document metadata and missing document.'
+    );
+
+    for (const q of this.queries) {
+      await saveNamedQuery(this.localStore, q);
+    }
+
+    const changedDocs = await applyBundleDocuments(
+      this.localStore,
+      this.documents
+    );
+
+    this.progress.taskState = 'Success';
+    return new BundleLoadResult({ ...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';
@@ -26,15 +27,15 @@ import { newDatastore } from '../remote/datastore';
 import { RemoteStore } from '../remote/remote_store';
 import { AsyncQueue, wrapInUserErrorIfRecoverable } from '../util/async_queue';
 import { Code, FirestoreError } from '../util/error';
-import { logDebug } from '../util/log';
+import { logDebug, logWarn } from '../util/log';
 import { Deferred } from '../util/promise';
 import {
   EventManager,
   ListenOptions,
   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,27 @@ 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(async () => {
+      loadBundle(this.syncEngine, reader, task);
+      return task.catch(e => {
+        logWarn(LOG_TAG, `Loading bundle failed with ${e}`);
+      });
+    });
+
+    return task;
+  }
 }