Add option to allow SDK create cache indexes automatically to improve query execution locally (#4987)

* Add counter

* address feedback 1

* add copyright

* fix concurrency bug

* implement autoClientIndexing

* Add tests and fix bugs for BuildTargetIndex

* hide getter from public API

* move the flag from IndexManager to QueryEngine

* Address feedback

* move auto index flag to runtime

* Support old way to enable persistent for PersistentCacheManager

* Polish Tests

* Add hide and copyright

* clean up unused function

* Rename PersistentCacheManager to PersistentCacheIndexManager

* Remove unused QueryContext

* Address feedbacks other than adding tests and comments

* Change the api to match the update

* Add tests

* Increase tests coverage

* Add comments

* add configurable min documents to create indexes

* Address Denver's feedback

* Address feedback

* address more feedbacks

* use the number getting from 100 ~ 1000 documents experiment

* Address feedbacks

* improve debug log

Author: cherylEnkidu

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/IndexingTest.java

@@ -14,7 +14,13 @@
 
 package com.google.firebase.firestore;
 
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollectionWithDocs;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testFirestore;
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.waitFor;
+import static com.google.firebase.firestore.testutil.TestUtil.assertDoesNotThrow;
+import static com.google.firebase.firestore.testutil.TestUtil.expectError;
+import static com.google.firebase.firestore.testutil.TestUtil.map;
+import static org.junit.Assert.assertEquals;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.android.gms.tasks.Task;
@@ -101,4 +107,73 @@ public void testBadIndexDoesNotCrashClient() {
                     + "  \"fieldOverrides\": []\n"
                     + "}"));
   }
+
+  @Test
+  public void testAutoIndexCreationSetSuccessfully() {
+    // Use persistent disk cache (default)
+    FirebaseFirestore db = testFirestore();
+    FirebaseFirestoreSettings settings =
+        new FirebaseFirestoreSettings.Builder(db.getFirestoreSettings())
+            .setLocalCacheSettings(PersistentCacheSettings.newBuilder().build())
+            .build();
+    db.setFirestoreSettings(settings);
+
+    CollectionReference collection =
+        testCollectionWithDocs(
+            map(
+                "a", map("match", true),
+                "b", map("match", false),
+                "c", map("match", false)));
+    QuerySnapshot results = waitFor(collection.whereEqualTo("match", true).get());
+    assertEquals(1, results.size());
+
+    assertDoesNotThrow(() -> db.getPersistentCacheIndexManager().enableIndexAutoCreation());
+
+    results = waitFor(collection.whereEqualTo("match", true).get());
+    assertEquals(1, results.size());
+
+    assertDoesNotThrow(() -> db.getPersistentCacheIndexManager().disableIndexAutoCreation());
+
+    results = waitFor(collection.whereEqualTo("match", true).get());
+    assertEquals(1, results.size());
+  }
+
+  @Test
+  public void testAutoIndexCreationSetSuccessfullyUsingDefault() {
+    // Use persistent disk cache (default)
+    FirebaseFirestore db = testFirestore();
+
+    CollectionReference collection =
+        testCollectionWithDocs(
+            map(
+                "a", map("match", true),
+                "b", map("match", false),
+                "c", map("match", false)));
+    QuerySnapshot results = waitFor(collection.whereEqualTo("match", true).get());
+    assertEquals(1, results.size());
+
+    assertDoesNotThrow(() -> db.getPersistentCacheIndexManager().enableIndexAutoCreation());
+
+    results = waitFor(collection.whereEqualTo("match", true).get());
+    assertEquals(1, results.size());
+
+    assertDoesNotThrow(() -> db.getPersistentCacheIndexManager().disableIndexAutoCreation());
+
+    results = waitFor(collection.whereEqualTo("match", true).get());
+    assertEquals(1, results.size());
+  }
+
+  @Test
+  public void testAutoIndexCreationAfterFailsTermination() {
+    FirebaseFirestore db = testFirestore();
+    waitFor(db.terminate());
+
+    expectError(
+        () -> db.getPersistentCacheIndexManager().enableIndexAutoCreation(),
+        "The client has already been terminated");
+
+    expectError(
+        () -> db.getPersistentCacheIndexManager().disableIndexAutoCreation(),
+        "The client has already been terminated");
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/FirebaseFirestore.java

@@ -23,6 +23,7 @@
 import androidx.annotation.Keep;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.RestrictTo;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.tasks.Task;
 import com.google.android.gms.tasks.TaskCompletionSource;
@@ -104,6 +105,8 @@ public interface InstanceRegistry {
   private volatile FirestoreClient client;
   private final GrpcMetadataProvider metadataProvider;
 
+  @Nullable private PersistentCacheIndexManager persistentCacheIndexManager;
+
   @NonNull
   private static FirebaseApp getDefaultFirebaseApp() {
     FirebaseApp app = FirebaseApp.getInstance();
@@ -403,6 +406,26 @@ public Task<Void> setIndexConfiguration(@NonNull String json) {
     return client.configureFieldIndexes(parsedIndexes);
   }
 
+  /**
+   * Returns the PersistentCache Index Manager used by this {@code FirebaseFirestore} object.
+   *
+   * @return The {@code PersistentCacheIndexManager} instance or null if local persistent storage is
+   *     not in use.
+   */
+  // TODO(csi): Remove the `hide` and scope annotations.
+  /** @hide */
+  @RestrictTo(RestrictTo.Scope.LIBRARY)
+  @Nullable
+  public synchronized PersistentCacheIndexManager getPersistentCacheIndexManager() {
+    ensureClientConfigured();
+    if (persistentCacheIndexManager == null
+        && (settings.isPersistenceEnabled()
+            || settings.getCacheSettings() instanceof PersistentCacheSettings)) {
+      persistentCacheIndexManager = new PersistentCacheIndexManager(client);
+    }
+    return persistentCacheIndexManager;
+  }
+
   /**
    * Gets a {@code CollectionReference} instance that refers to the collection at the specified path
    * within the database.

firebase-firestore/src/main/java/com/google/firebase/firestore/PersistentCacheIndexManager.java

@@ -0,0 +1,54 @@
+// 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.
+
+package com.google.firebase.firestore;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.RestrictTo;
+import com.google.firebase.firestore.core.FirestoreClient;
+
+/**
+ * A {@code PersistentCacheIndexManager} which you can config persistent cache indexes used for
+ * local query execution.
+ *
+ * <p>To use, call {@link FirebaseFirestore#getPersistentCacheIndexManager()} to get an instance.
+ */
+// TODO(csi): Remove the `hide` and scope annotations.
+/** @hide */
+@RestrictTo(RestrictTo.Scope.LIBRARY)
+public final class PersistentCacheIndexManager {
+  @NonNull private FirestoreClient client;
+
+  PersistentCacheIndexManager(FirestoreClient client) {
+    this.client = client;
+  }
+
+  /**
+   * Enables SDK to create persistent cache indexes automatically for local query execution when SDK
+   * believes cache indexes can help improves performance.
+   *
+   * <p>This feature is disabled by default.
+   */
+  public void enableIndexAutoCreation() {
+    client.setIndexAutoCreationEnabled(true);
+  }
+
+  /**
+   * Stops creating persistent cache indexes automatically for local query execution. The indexes
+   * which have been created by calling {@link #enableIndexAutoCreation()} still take effect.
+   */
+  public void disableIndexAutoCreation() {
+    client.setIndexAutoCreationEnabled(false);
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/core/FirestoreClient.java

@@ -353,6 +353,11 @@ public Task<Void> configureFieldIndexes(List<FieldIndex> fieldIndices) {
     return asyncQueue.enqueue(() -> localStore.configureFieldIndexes(fieldIndices));
   }
 
+  public void setIndexAutoCreationEnabled(boolean enabled) {
+    verifyNotTerminated();
+    asyncQueue.enqueueAndForget(() -> localStore.setIndexAutoCreationEnabled(enabled));
+  }
+
   public void removeSnapshotsInSyncListener(EventListener<Void> listener) {
     // Checks for shutdown but does not raise error, allowing remove after shutdown to be a no-op.
     if (isTerminated()) {

firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexManager.java

@@ -78,6 +78,9 @@ enum IndexType {
   /** Removes the given field index and deletes all index values. */
   void deleteFieldIndex(FieldIndex index);
 
+  /** Creates a full matched field index which serves the given target. */
+  void createTargetIndexes(Target target);
+
   /**
    * Returns a list of field indexes that correspond to the specified collection group.
    *

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -258,19 +258,32 @@ void recalculateAndSaveOverlays(Set<DocumentKey> documentKeys) {
    *
    * @param query The query to match documents against.
    * @param offset Read time and key to start scanning by (exclusive).
+   * @param context A optional tracker to keep a record of important details during database local
+   *     query execution.
    */
   ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
-      Query query, IndexOffset offset) {
+      Query query, IndexOffset offset, @Nullable QueryContext context) {
     ResourcePath path = query.getPath();
     if (query.isDocumentQuery()) {
       return getDocumentsMatchingDocumentQuery(path);
     } else if (query.isCollectionGroupQuery()) {
-      return getDocumentsMatchingCollectionGroupQuery(query, offset);
+      return getDocumentsMatchingCollectionGroupQuery(query, offset, context);
     } else {
-      return getDocumentsMatchingCollectionQuery(query, offset);
+      return getDocumentsMatchingCollectionQuery(query, offset, context);
     }
   }
 
+  /**
+   * Performs a query against the local view of all documents.
+   *
+   * @param query The query to match documents against.
+   * @param offset Read time and key to start scanning by (exclusive).
+   */
+  ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
+      Query query, IndexOffset offset) {
+    return getDocumentsMatchingQuery(query, offset, /*context*/ null);
+  }
+
   /** Performs a simple document lookup for the given path. */
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingDocumentQuery(
       ResourcePath path) {
@@ -284,7 +297,7 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingDocumentQu
   }
 
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollectionGroupQuery(
-      Query query, IndexOffset offset) {
+      Query query, IndexOffset offset, @Nullable QueryContext context) {
     hardAssert(
         query.getPath().isEmpty(),
         "Currently we only support collection group queries at the root.");
@@ -297,7 +310,7 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
     for (ResourcePath parent : parents) {
       Query collectionQuery = query.asCollectionQueryAtPath(parent.append(collectionId));
       ImmutableSortedMap<DocumentKey, Document> collectionResults =
-          getDocumentsMatchingCollectionQuery(collectionQuery, offset);
+          getDocumentsMatchingCollectionQuery(collectionQuery, offset, context);
       for (Map.Entry<DocumentKey, Document> docEntry : collectionResults) {
         results = results.insert(docEntry.getKey(), docEntry.getValue());
       }
@@ -362,11 +375,11 @@ private void populateOverlays(Map<DocumentKey, Overlay> overlays, Set<DocumentKe
   }
 
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollectionQuery(
-      Query query, IndexOffset offset) {
+      Query query, IndexOffset offset, @Nullable QueryContext context) {
     Map<DocumentKey, Overlay> overlays =
         documentOverlayCache.getOverlays(query.getPath(), offset.getLargestBatchId());
     Map<DocumentKey, MutableDocument> remoteDocuments =
-        remoteDocumentCache.getDocumentsMatchingQuery(query, offset, overlays.keySet());
+        remoteDocumentCache.getDocumentsMatchingQuery(query, offset, overlays.keySet(), context);
 
     // As documents might match the query because of their overlay we need to include documents
     // for all overlays in the initial document set.

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalStore.java

@@ -802,6 +802,10 @@ public void configureFieldIndexes(List<FieldIndex> newFieldIndexes) {
         });
   }
 
+  public void setIndexAutoCreationEnabled(boolean enabled) {
+    queryEngine.setIndexAutoCreationEnabled(enabled);
+  }
+
   /** Mutable state for the transaction in allocateQuery. */
   private static class AllocateQueryHolder {
     TargetData cached;

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryIndexManager.java

@@ -60,6 +60,9 @@ public void deleteFieldIndex(FieldIndex index) {
     // Field indices are not supported with memory persistence.
   }
 
+  @Override
+  public void createTargetIndexes(Target target) {}
+
   @Override
   @Nullable
   public List<DocumentKey> getDocumentsMatchingTarget(Target target) {

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java

@@ -18,6 +18,7 @@
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
@@ -97,7 +98,10 @@ public Map<DocumentKey, MutableDocument> getAll(
 
   @Override
   public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
-      Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
+      Query query,
+      IndexOffset offset,
+      @Nonnull Set<DocumentKey> mutatedKeys,
+      @Nullable QueryContext context) {
     Map<DocumentKey, MutableDocument> result = new HashMap<>();
 
     // Documents are ordered by key, so we can use a prefix scan to narrow down the documents
@@ -135,6 +139,12 @@ public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
     return result;
   }
 
+  @Override
+  public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
+      Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
+    return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
+  }
+
   Iterable<Document> getDocuments() {
     return new DocumentIterable();
   }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/QueryContext.java

@@ -0,0 +1,29 @@
+// 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.
+
+package com.google.firebase.firestore.local;
+
+/** A tracker to keep a record of important details during database local query execution. */
+public class QueryContext {
+  /** Counts the number of documents passed through during local query execution. */
+  private int documentReadCount = 0;
+
+  public int getDocumentReadCount() {
+    return documentReadCount;
+  }
+
+  public void incrementDocumentReadCount() {
+    documentReadCount++;
+  }
+}