Keep track of number of queries in the query cache

packages/firestore/src/local/indexeddb_persistence.ts

@@ -25,8 +25,13 @@ import { AutoId } from '../util/misc';
 import { IndexedDbMutationQueue } from './indexeddb_mutation_queue';
 import { IndexedDbQueryCache } from './indexeddb_query_cache';
 import { IndexedDbRemoteDocumentCache } from './indexeddb_remote_document_cache';
-import { ALL_STORES, DbOwner, DbOwnerKey } from './indexeddb_schema';
-import { createOrUpgradeDb, SCHEMA_VERSION } from './indexeddb_schema';
+import {
+  ALL_STORES,
+  createOrUpgradeDb,
+  DbOwner,
+  DbOwnerKey,
+  SCHEMA_VERSION
+} from './indexeddb_schema';
 import { LocalSerializer } from './local_serializer';
 import { MutationQueue } from './mutation_queue';
 import { Persistence } from './persistence';

packages/firestore/src/local/indexeddb_query_cache.ts

@@ -31,8 +31,7 @@ import {
   DbTargetDocumentKey,
   DbTargetGlobal,
   DbTargetGlobalKey,
-  DbTargetKey,
-  DbTimestamp
+  DbTargetKey
 } from './indexeddb_schema';
 import { LocalSerializer } from './local_serializer';
 import { PersistenceTransaction } from './persistence';
@@ -56,7 +55,8 @@ export class IndexedDbQueryCache implements QueryCache {
   private metadata = new DbTargetGlobal(
     /*highestTargetId=*/ 0,
     /*lastListenSequenceNumber=*/ 0,
-    SnapshotVersion.MIN.toTimestamp()
+    SnapshotVersion.MIN.toTimestamp(),
+    /*targetCount=*/ 0
   );
 
   /** The garbage collector to notify about potential garbage keys. */
@@ -97,36 +97,75 @@ export class IndexedDbQueryCache implements QueryCache {
     );
   }
 
-  addQueryData(
+  private saveQueryData(
     transaction: PersistenceTransaction,
     queryData: QueryData
   ): PersistencePromise<void> {
-    const targetId = queryData.targetId;
-    const addedQueryPromise = targetsStore(transaction).put(
-      this.serializer.toDbTarget(queryData)
-    );
-    if (targetId > this.metadata.highestTargetId) {
-      this.metadata.highestTargetId = targetId;
-      return addedQueryPromise.next(() =>
-        globalTargetStore(transaction).put(DbTargetGlobal.key, this.metadata)
-      );
-    } else {
-      return addedQueryPromise;
+    return targetsStore(transaction).put(this.serializer.toDbTarget(queryData));
+  }
+
+  // Updates the in-memory version of the metadata. Saving is done separately.
+  // Returns true if there were any changes to the metadata.
+  private updateMetadata(queryData: QueryData): boolean {
+    let needsUpdate = false;
+    if (queryData.targetId > this.metadata.highestTargetId) {
+      this.metadata.highestTargetId = queryData.targetId;
+      needsUpdate = true;
     }
+
+    // TODO(gsoltis): add sequence number check
+    return needsUpdate;
   }
 
-  removeQueryData(
+  private saveMetadata(
+    transaction: PersistenceTransaction
+  ): PersistencePromise<void> {
+    return globalTargetStore(transaction).put(
+      DbTargetGlobal.key,
+      this.metadata
+    );
+  }
+
+  addQueryData(
+    transaction: PersistenceTransaction,
+    queryData: QueryData
+  ): PersistencePromise<void> {
+    return this.saveQueryData(transaction, queryData).next(() => {
+      this.metadata.targetCount += 1;
+      this.updateMetadata(queryData);
+      return this.saveMetadata(transaction);
+    });
+  }
+
+  updateQueryData(
     transaction: PersistenceTransaction,
     queryData: QueryData
   ): PersistencePromise<void> {
-    return this.removeMatchingKeysForTargetId(
-      transaction,
-      queryData.targetId
-    ).next(() => {
-      targetsStore(transaction).delete(queryData.targetId);
+    return this.saveQueryData(transaction, queryData).next(() => {
+      if (this.updateMetadata(queryData)) {
+        return this.saveMetadata(transaction);
+      } else {
+        return PersistencePromise.resolve();
+      }
     });
   }
 
+  removeQueryData(
+    transaction: PersistenceTransaction,
+    queryData: QueryData
+  ): PersistencePromise<void> {
+    return this.removeMatchingKeysForTargetId(transaction, queryData.targetId)
+      .next(() => targetsStore(transaction).delete(queryData.targetId))
+      .next(() => {
+        this.metadata.targetCount -= 1;
+        return this.saveMetadata(transaction);
+      });
+  }
+
+  count(): number {
+    return this.metadata.targetCount;
+  }
+
   getQueryData(
     transaction: PersistenceTransaction,
     query: Query

packages/firestore/src/local/indexeddb_schema.ts

@@ -21,54 +21,55 @@ import { ResourcePath } from '../model/path';
 import { assert } from '../util/assert';
 
 import { encode, EncodedResourcePath } from './encoded_resource_path';
+import { wrapRequest } from './simple_db';
+import { PersistencePromise } from './persistence_promise';
+import { SnapshotVersion } from '../core/snapshot_version';
 
-export const SCHEMA_VERSION = 1;
-
-/** Performs database creation and (in the future) upgrades between versions. */
-export function createOrUpgradeDb(db: IDBDatabase, oldVersion: number): void {
-  assert(oldVersion === 0, 'Unexpected upgrade from version ' + oldVersion);
-
-  db.createObjectStore(DbMutationQueue.store, {
-    keyPath: DbMutationQueue.keyPath
-  });
-
-  // TODO(mikelehen): Get rid of "as any" if/when TypeScript fixes their
-  // types. https://github.com/Microsoft/TypeScript/issues/14322
-  db.createObjectStore(
-    DbMutationBatch.store,
-    // tslint:disable-next-line:no-any
-    { keyPath: DbMutationBatch.keyPath as any }
-  );
+/**
+ * Schema Version for the Web client (containing the Mutation Queue, the Query
+ * and the Remote Document Cache) and target count added to the target global.
+ */
+export const SCHEMA_VERSION = 2;
 
-  const targetDocumentsStore = db.createObjectStore(
-    DbTargetDocument.store,
-    // tslint:disable-next-line:no-any
-    { keyPath: DbTargetDocument.keyPath as any }
-  );
-  targetDocumentsStore.createIndex(
-    DbTargetDocument.documentTargetsIndex,
-    DbTargetDocument.documentTargetsKeyPath,
-    { unique: true }
+/**
+ * Performs database creation and schema upgrades.
+ *
+ * Note that in production, this method is only ever used to upgrade the schema
+ * to SCHEMA_VERSION. Different versions are only used for testing and
+ * local feature development.
+ */
+export function createOrUpgradeDb(
+  db: IDBDatabase,
+  txn: IDBTransaction,
+  fromVersion: number,
+  toVersion: number
+): PersistencePromise<void> {
+  // This function currently supports migrating to schema version 1 (Mutation
+  // Queue, Query and Remote Document Cache) and schema version 2 (Query
+  // counting).
+  assert(
+    fromVersion < toVersion && fromVersion >= 0 && toVersion <= 2,
+    'Unexpected schema upgrade from v${fromVersion} to v{toVersion}.'
   );
 
-  const targetStore = db.createObjectStore(DbTarget.store, {
-    keyPath: DbTarget.keyPath
-  });
-  // NOTE: This is unique only because the TargetId is the suffix.
-  targetStore.createIndex(
-    DbTarget.queryTargetsIndexName,
-    DbTarget.queryTargetsKeyPath,
-    { unique: true }
-  );
+  if (fromVersion < 1 && toVersion >= 1) {
+    createOwnerStore(db);
+    createMutationQueue(db);
+    createQueryCache(db);
+    createRemoteDocumentCache(db);
+  }
 
-  // NOTE: keys for these stores are specified explicitly rather than using a
-  // keyPath.
-  db.createObjectStore(DbDocumentMutation.store);
-  db.createObjectStore(DbRemoteDocument.store);
-  db.createObjectStore(DbOwner.store);
-  db.createObjectStore(DbTargetGlobal.store);
+  let p = PersistencePromise.resolve();
+  if (fromVersion < 2 && toVersion >= 2) {
+    p = p.next(() => ensureTargetGlobal(txn)).next(() => saveTargetCount(txn));
+  }
+  return p;
 }
 
+// TODO(mikelehen): Get rid of "as any" if/when TypeScript fixes their types.
+// https://github.com/Microsoft/TypeScript/issues/14322
+type KeyPath = any; // tslint:disable-line:no-any
+
 /**
  * Wrapper class to store timestamps (seconds and nanos) in IndexedDb objects.
  */
@@ -94,6 +95,10 @@ export class DbOwner {
   constructor(public ownerId: string, public leaseTimestampMs: number) {}
 }
 
+function createOwnerStore(db: IDBDatabase): void {
+  db.createObjectStore(DbOwner.store);
+}
+
 /** Object keys in the 'mutationQueues' store are userId strings. */
 export type DbMutationQueueKey = string;
 
@@ -183,6 +188,18 @@ export class DbMutationBatch {
  */
 export type DbDocumentMutationKey = [string, EncodedResourcePath, BatchId];
 
+function createMutationQueue(db: IDBDatabase): void {
+  db.createObjectStore(DbMutationQueue.store, {
+    keyPath: DbMutationQueue.keyPath
+  });
+
+  db.createObjectStore(DbMutationBatch.store, {
+    keyPath: DbMutationBatch.keyPath as KeyPath
+  });
+
+  db.createObjectStore(DbDocumentMutation.store);
+}
+
 /**
  * An object to be stored in the 'documentMutations' store in IndexedDb.
  *
@@ -241,6 +258,10 @@ export class DbDocumentMutation {
  */
 export type DbRemoteDocumentKey = string[];
 
+function createRemoteDocumentCache(db: IDBDatabase): void {
+  db.createObjectStore(DbRemoteDocument.store);
+}
+
 /**
  * Represents the known absence of a document at a particular version.
  * Stored in IndexedDb as part of a DbRemoteDocument object.
@@ -451,15 +472,91 @@ export class DbTargetGlobal {
      * until the backend has caught up to this snapshot version again. This
      * prevents our cache from ever going backwards in time.
      */
-    public lastRemoteSnapshotVersion: DbTimestamp
+    public lastRemoteSnapshotVersion: DbTimestamp,
+    /**
+     * The number of targets persisted.
+     */
+    public targetCount: number
   ) {}
 }
 
+function createQueryCache(db: IDBDatabase): void {
+  const targetDocumentsStore = db.createObjectStore(DbTargetDocument.store, {
+    keyPath: DbTargetDocument.keyPath as KeyPath
+  });
+  targetDocumentsStore.createIndex(
+    DbTargetDocument.documentTargetsIndex,
+    DbTargetDocument.documentTargetsKeyPath,
+    { unique: true }
+  );
+
+  const targetStore = db.createObjectStore(DbTarget.store, {
+    keyPath: DbTarget.keyPath
+  });
+
+  // NOTE: This is unique only because the TargetId is the suffix.
+  targetStore.createIndex(
+    DbTarget.queryTargetsIndexName,
+    DbTarget.queryTargetsKeyPath,
+    { unique: true }
+  );
+  db.createObjectStore(DbTargetGlobal.store);
+}
+
+/**
+ * Counts the number of targets persisted and adds that value to the target
+ * global singleton.
+ * @param {IDBTransaction} txn The version upgrade transaction for indexeddb
+ * @return {PersistencePromise<void>} A promise of the operations to add the
+ *                                    count to the metadata row.
+ */
+function saveTargetCount(txn: IDBTransaction): PersistencePromise<void> {
+  const globalStore = txn.objectStore(DbTargetGlobal.store);
+  return wrapRequest<number>(txn.objectStore(DbTarget.store).count()).next(
+    (count: number) =>
+      wrapRequest<DbTargetGlobal>(globalStore.get(DbTargetGlobal.key)).next(
+        (metadata: DbTargetGlobal) => {
+          metadata.targetCount = count;
+          return wrapRequest<void>(
+            globalStore.put(metadata, DbTargetGlobal.key)
+          );
+        }
+      )
+  );
+}
+
 /**
- * The list of all IndexedDB stored used by the SDK. This is used when creating
- * transactions so that access across all stores is done atomically.
+ * Ensures that the target global singleton row exists by adding it if it's
+ * missing.
+ * @param {IDBTransaction} txn The version upgrade transaction for indexeddb
+ * @return {PersistencePromise<void>} A promise of operations to ensure that the
+ *                                    metadata row exists
  */
-export const ALL_STORES = [
+function ensureTargetGlobal(txn: IDBTransaction): PersistencePromise<void> {
+  const globalStore = txn.objectStore(DbTargetGlobal.store);
+  return wrapRequest<DbTargetGlobal | null>(
+    globalStore.get(DbTargetGlobal.key)
+  ).next((metadata: DbTargetGlobal | null) => {
+    if (metadata != null) {
+      return PersistencePromise.resolve();
+    } else {
+      return wrapRequest<void>(
+        globalStore.put(
+          new DbTargetGlobal(
+            /*highestTargetId=*/ 0,
+            /*lastListenSequenceNumber=*/ 0,
+            SnapshotVersion.MIN.toTimestamp(),
+            /*targetCount=*/ 0
+          ),
+          DbTargetGlobal.key
+        )
+      );
+    }
+  });
+}
+
+// Visible for testing
+export const V1_STORES = [
   DbMutationQueue.store,
   DbMutationBatch.store,
   DbDocumentMutation.store,
@@ -469,3 +566,10 @@ export const ALL_STORES = [
   DbTargetGlobal.store,
   DbTargetDocument.store
 ];
+
+/**
+ * The list of all default IndexedDB stores used throughout the SDK. This is
+ * used when creating transactions so that access across all stores is done
+ * atomically.
+ */
+export const ALL_STORES = [...V1_STORES];

packages/firestore/src/local/local_store.ts

@@ -478,7 +478,7 @@ export class LocalStore {
               snapshotVersion: change.snapshotVersion
             });
             this.targetIds[targetId] = queryData;
-            promises.push(this.queryCache.addQueryData(txn, queryData));
+            promises.push(this.queryCache.updateQueryData(txn, queryData));
           }
         }
       );