Firestore: enable/disablePersistentCacheIndexAutoCreation() added, but hidden for now

packages/firestore/src/api.ts

@@ -223,6 +223,12 @@ export type {
 } from './api/credentials';
 export { EmptyAuthCredentialsProvider as _EmptyAuthCredentialsProvider } from './api/credentials';
 export { EmptyAppCheckTokenProvider as _EmptyAppCheckTokenProvider } from './api/credentials';
+export {
+  PersistentCacheIndexManager as _PersistentCacheIndexManager,
+  getPersistentCacheIndexManager as _getPersistentCacheIndexManager,
+  enablePersistentCacheIndexAutoCreation as _enablePersistentCacheIndexAutoCreation,
+  disablePersistentCacheIndexAutoCreation as _disablePersistentCacheIndexAutoCreation
+} from './api/persistent_cache_index_manager';
 export {
   ExistenceFilterMismatchCallback as _TestingHooksExistenceFilterMismatchCallback,
   TestingHooks as _TestingHooks

packages/firestore/src/api/persistent_cache_index_manager.ts

@@ -0,0 +1,166 @@
+/**
+ * @license
+ * Copyright 2023 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 {
+  firestoreClientSetPersistentCacheIndexAutoCreationEnabled,
+  FirestoreClient,
+  TestingHooks as FirestoreClientTestingHooks
+} from '../core/firestore_client';
+import { cast } from '../util/input_validation';
+import { logDebug, logWarn } from '../util/log';
+import { testingHooksSpi } from '../util/testing_hooks_spi';
+
+import { ensureFirestoreConfigured, Firestore } from './database';
+
+/**
+ * A `PersistentCacheIndexManager` which you can config persistent cache indexes
+ * used for local query execution.
+ *
+ * To use, call `getPersistentCacheIndexManager()` to get an instance.
+ *
+ * TODO(CSI) Remove @internal to make the API publicly available.
+ * @internal
+ */
+export class PersistentCacheIndexManager {
+  readonly type: 'PersistentCacheIndexManager' = 'PersistentCacheIndexManager';
+
+  /** @hideconstructor */
+  constructor(readonly _client: FirestoreClient) {}
+}
+
+/**
+ * Returns the PersistentCache Index Manager used by the given `Firestore`
+ * object.
+ *
+ * @return The `PersistentCacheIndexManager` instance, or `null` if local
+ * persistent storage is not in use.
+ *
+ * TODO(CSI) Remove @internal to make the API publicly available.
+ * @internal
+ */
+export function getPersistentCacheIndexManager(
+  firestore: Firestore
+): PersistentCacheIndexManager | null {
+  firestore = cast(firestore, Firestore);
+
+  const cachedInstance = persistentCacheIndexManagerByFirestore.get(firestore);
+  if (cachedInstance) {
+    return cachedInstance;
+  }
+
+  const client = ensureFirestoreConfigured(firestore);
+  if (client._uninitializedComponentsProvider?._offlineKind !== 'persistent') {
+    return null;
+  }
+
+  const instance = new PersistentCacheIndexManager(client);
+  persistentCacheIndexManagerByFirestore.set(firestore, instance);
+  return instance;
+}
+
+/**
+ * Enables SDK to create persistent cache indexes automatically for local query
+ * execution when SDK believes cache indexes can help improves performance.
+ *
+ * This feature is disabled by default.
+ *
+ * TODO(CSI) Remove @internal to make the API publicly available.
+ * @internal
+ */
+export function enablePersistentCacheIndexAutoCreation(
+  indexManager: PersistentCacheIndexManager
+): void {
+  setPersistentCacheIndexAutoCreationEnabled(indexManager, true);
+}
+
+/**
+ * Stops creating persistent cache indexes automatically for local query
+ * execution. The indexes which have been created by calling
+ * `enablePersistentCacheIndexAutoCreation()` still take effect.
+ *
+ * TODO(CSI) Remove @internal to make the API publicly available.
+ * @internal
+ */
+export function disablePersistentCacheIndexAutoCreation(
+  indexManager: PersistentCacheIndexManager
+): void {
+  setPersistentCacheIndexAutoCreationEnabled(indexManager, false);
+}
+
+function setPersistentCacheIndexAutoCreationEnabled(
+  indexManager: PersistentCacheIndexManager,
+  isEnabled: boolean
+): void {
+  indexManager._client.verifyNotTerminated();
+
+  const promise = firestoreClientSetPersistentCacheIndexAutoCreationEnabled(
+    indexManager._client,
+    isEnabled
+  );
+
+  promise
+    .then(_ =>
+      logDebug(
+        `setting persistent cache index auto creation ` +
+          `isEnabled=${isEnabled} succeeded`
+      )
+    )
+    .catch(error =>
+      logWarn(
+        `setting persistent cache index auto creation ` +
+          `isEnabled=${isEnabled} failed`,
+        error
+      )
+    );
+
+  testingHooksSpi?.notifyPersistentCacheIndexAutoCreationToggle(promise);
+}
+
+/**
+ * Maps `Firestore` instances to their corresponding
+ * `PersistentCacheIndexManager` instances.
+ *
+ * Use a `WeakMap` so that the mapping will be automatically dropped when the
+ * `Firestore` instance is garbage collected. This emulates a private member
+ * as described in https://goo.gle/454yvug.
+ */
+const persistentCacheIndexManagerByFirestore = new WeakMap<
+  Firestore,
+  PersistentCacheIndexManager
+>();
+
+/**
+ * Test-only hooks into the SDK for use exclusively by integration tests.
+ */
+export class TestingHooks {
+  private constructor() {
+    throw new Error('creating instances is not supported');
+  }
+
+  static setIndexAutoCreationSettings(
+    indexManager: PersistentCacheIndexManager,
+    settings: {
+      indexAutoCreationMinCollectionSize?: number;
+      relativeIndexReadCostPerDocument?: number;
+    }
+  ): Promise<void> {
+    return FirestoreClientTestingHooks.setPersistentCacheIndexAutoCreationSettings(
+      indexManager._client,
+      settings
+    );
+  }
+}

packages/firestore/src/core/firestore_client.ts

@@ -23,13 +23,16 @@ import {
   CredentialsProvider
 } from '../api/credentials';
 import { User } from '../auth/user';
+import { IndexType } from '../local/index_manager';
 import { LocalStore } from '../local/local_store';
 import {
   localStoreConfigureFieldIndexes,
   localStoreExecuteQuery,
   localStoreGetNamedQuery,
   localStoreHandleUserChange,
-  localStoreReadDocument
+  localStoreReadDocument,
+  localStoreSetIndexAutoCreationEnabled,
+  TestingHooks as LocalStoreTestingHooks
 } from '../local/local_store_impl';
 import { Persistence } from '../local/persistence';
 import { Document } from '../model/document';
@@ -828,3 +831,51 @@ export function firestoreClientSetIndexConfiguration(
     );
   });
 }
+
+export function firestoreClientSetPersistentCacheIndexAutoCreationEnabled(
+  client: FirestoreClient,
+  isEnabled: boolean
+): Promise<void> {
+  return client.asyncQueue.enqueue(async () => {
+    return localStoreSetIndexAutoCreationEnabled(
+      await getLocalStore(client),
+      isEnabled
+    );
+  });
+}
+
+/**
+ * Test-only hooks into the SDK for use exclusively by integration tests.
+ */
+export class TestingHooks {
+  private constructor() {
+    throw new Error('creating instances is not supported');
+  }
+
+  static getQueryIndexType(
+    client: FirestoreClient,
+    query: Query
+  ): Promise<IndexType> {
+    return client.asyncQueue.enqueue(async () => {
+      const localStore = await getLocalStore(client);
+      return LocalStoreTestingHooks.getQueryIndexType(localStore, query);
+    });
+  }
+
+  static setPersistentCacheIndexAutoCreationSettings(
+    client: FirestoreClient,
+    settings: {
+      indexAutoCreationMinCollectionSize?: number;
+      relativeIndexReadCostPerDocument?: number;
+    }
+  ): Promise<void> {
+    const settingsCopy = { ...settings };
+    return client.asyncQueue.enqueue(async () => {
+      const localStore = await getLocalStore(client);
+      LocalStoreTestingHooks.setIndexAutoCreationSettings(
+        localStore,
+        settingsCopy
+      );
+    });
+  }
+}

packages/firestore/src/local/index_manager.ts

@@ -41,6 +41,19 @@ export const enum IndexType {
   FULL
 }
 
+export function displayNameForIndexType(indexType: IndexType): string {
+  switch (indexType) {
+    case IndexType.NONE:
+      return 'NONE';
+    case IndexType.PARTIAL:
+      return 'PARTIAL';
+    case IndexType.FULL:
+      return 'FULL';
+    default:
+      return `[unknown IndexType: ${indexType}]`;
+  }
+}
+
 /**
  * Represents a set of indexes that are used to execute queries efficiently.
  *
@@ -92,6 +105,12 @@ export interface IndexManager {
     index: FieldIndex
   ): PersistencePromise<void>;
 
+  /** Creates a full matched field index which serves the given target. */
+  createTargetIndexes(
+    transaction: PersistenceTransaction,
+    target: Target
+  ): PersistencePromise<void>;
+
   /**
    * Returns a list of field indexes that correspond to the specified collection
    * group.

packages/firestore/src/local/indexeddb_index_manager.ts

@@ -252,6 +252,26 @@ export class IndexedDbIndexManager implements IndexManager {
       );
   }
 
+  createTargetIndexes(
+    transaction: PersistenceTransaction,
+    target: Target
+  ): PersistencePromise<void> {
+    return PersistencePromise.forEach(
+      this.getSubTargets(target),
+      (subTarget: Target) => {
+        return this.getIndexType(transaction, subTarget).next(type => {
+          if (type === IndexType.NONE || type === IndexType.PARTIAL) {
+            const targetIndexMatcher = new TargetIndexMatcher(subTarget);
+            return this.addFieldIndex(
+              transaction,
+              targetIndexMatcher.buildTargetIndex()
+            );
+          }
+        });
+      }
+    );
+  }
+
   getDocumentsMatchingTarget(
     transaction: PersistenceTransaction,
     target: Target

packages/firestore/src/local/indexeddb_remote_document_cache.ts

@@ -55,6 +55,7 @@ import {
 } from './local_serializer';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
+import { QueryContext } from './query_context';
 import { RemoteDocumentCache } from './remote_document_cache';
 import { RemoteDocumentChangeBuffer } from './remote_document_change_buffer';
 import { SimpleDbStore } from './simple_db';
@@ -279,7 +280,8 @@ class IndexedDbRemoteDocumentCacheImpl implements IndexedDbRemoteDocumentCache {
     transaction: PersistenceTransaction,
     query: Query,
     offset: IndexOffset,
-    mutatedDocs: OverlayMap
+    mutatedDocs: OverlayMap,
+    context?: QueryContext
   ): PersistencePromise<MutableDocumentMap> {
     const collection = query.path;
     const startKey = [
@@ -300,6 +302,7 @@ class IndexedDbRemoteDocumentCacheImpl implements IndexedDbRemoteDocumentCache {
     return remoteDocumentsStore(transaction)
       .loadAll(IDBKeyRange.bound(startKey, endKey, true))
       .next(dbRemoteDocs => {
+        context?.incrementDocumentReadCount(dbRemoteDocs.length);
         let results = mutableDocumentMap();
         for (const dbRemoteDoc of dbRemoteDocs) {
           const document = this.maybeDecodeDocument(