Add option to allow SDK create cache indexes automatically to improve query execution locally

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

@@ -18,6 +18,7 @@
 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;
 
@@ -108,7 +109,7 @@ public void testBadIndexDoesNotCrashClient() {
   }
 
   @Test
-  public void testAutomaticIndexingSetSuccessfully() {
+  public void testAutoIndexCreationSetSuccessfully() {
     // Use persistent disk cache (default)
     FirebaseFirestore db = testFirestore();
     FirebaseFirestoreSettings settings =
@@ -138,7 +139,7 @@ public void testAutomaticIndexingSetSuccessfully() {
   }
 
   @Test
-  public void testAutomaticIndexingSetSuccessfullyUseDefault() {
+  public void testAutoIndexCreationSetSuccessfullyUsingDefault() {
     // Use persistent disk cache (default)
     FirebaseFirestore db = testFirestore();
 
@@ -161,4 +162,18 @@ public void testAutomaticIndexingSetSuccessfullyUseDefault() {
     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

@@ -49,7 +49,6 @@
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.remote.FirestoreChannel;
 import com.google.firebase.firestore.remote.GrpcMetadataProvider;
-import com.google.firebase.firestore.util.Assert;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.ByteBufferInputStream;
 import com.google.firebase.firestore.util.Executors;
@@ -107,8 +106,7 @@ public interface InstanceRegistry {
   private volatile FirestoreClient client;
   private final GrpcMetadataProvider metadataProvider;
 
-  @Nullable
-  private AtomicReference<PersistentCacheIndexManager> persistentCacheIndexManager =
+  private final AtomicReference<PersistentCacheIndexManager> persistentCacheIndexManager =
       new AtomicReference<>();
 
   @NonNull
@@ -422,12 +420,11 @@ public Task<Void> setIndexConfiguration(@NonNull String json) {
   @Nullable
   public PersistentCacheIndexManager getPersistentCacheIndexManager() {
     ensureClientConfigured();
-    Assert.hardAssertNonNull(
-        persistentCacheIndexManager, "persistentCacheIndexManager has not been initialized.");
 
-    if (settings.isPersistenceEnabled()
-        || settings.getCacheSettings() instanceof PersistentCacheSettings) {
-      persistentCacheIndexManager.compareAndSet(null, new PersistentCacheIndexManager(client));
+    if ((settings.isPersistenceEnabled()
+            || settings.getCacheSettings() instanceof PersistentCacheSettings)
+        && persistentCacheIndexManager.get() == null) {
+      persistentCacheIndexManager.set(new PersistentCacheIndexManager(client));
     }
 
     return persistentCacheIndexManager.get();

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

@@ -22,16 +22,15 @@
  * A {@code PersistentCacheIndexManager} which you can config persistent cache indexes used for
  * local query execution.
  *
- * <p>To use, calling {@link FirebaseFirestore#getPersistentCacheIndexManager()} to get an instance.
+ * <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;
 
-  @RestrictTo(RestrictTo.Scope.LIBRARY)
-  public PersistentCacheIndexManager(FirestoreClient client) {
+  PersistentCacheIndexManager(FirestoreClient client) {
     this.client = client;
   }
 
@@ -42,14 +41,14 @@ public PersistentCacheIndexManager(FirestoreClient client) {
    * <p>This feature is disabled by default.
    */
   public void enableIndexAutoCreation() {
-    client.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.disableIndexAutoCreation();
+    client.setIndexAutoCreationEnabled(false);
   }
 }

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

@@ -353,14 +353,9 @@ public Task<Void> configureFieldIndexes(List<FieldIndex> fieldIndices) {
     return asyncQueue.enqueue(() -> localStore.configureFieldIndexes(fieldIndices));
   }
 
-  public void enableIndexAutoCreation() {
+  public void setIndexAutoCreationEnabled(boolean enabled) {
     verifyNotTerminated();
-    asyncQueue.enqueueAndForget(() -> localStore.enableIndexAutoCreation());
-  }
-
-  public void disableIndexAutoCreation() {
-    verifyNotTerminated();
-    asyncQueue.enqueueAndForget(() -> localStore.disableIndexAutoCreation());
+    asyncQueue.enqueueAndForget(() -> localStore.setIndexAutoCreationEnabled(enabled));
   }
 
   public void removeSnapshotsInSyncListener(EventListener<Void> listener) {

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

@@ -281,7 +281,7 @@ ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
    */
   ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       Query query, IndexOffset offset) {
-    return getDocumentsMatchingQuery(query, offset, null);
+    return getDocumentsMatchingQuery(query, offset, /*context*/ null);
   }
 
   /** Performs a simple document lookup for the given path. */

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

@@ -802,12 +802,8 @@ public void configureFieldIndexes(List<FieldIndex> newFieldIndexes) {
         });
   }
 
-  public void enableIndexAutoCreation() {
-    queryEngine.enableIndexAutoCreation();
-  }
-
-  public void disableIndexAutoCreation() {
-    queryEngine.disableIndexAutoCreation();
+  public void setIndexAutoCreationEnabled(boolean enabled) {
+    queryEngine.setIndexAutoCreationEnabled(enabled);
   }
 
   /** Mutable state for the transaction in allocateQuery. */

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

@@ -142,7 +142,7 @@ public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
   @Override
   public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
       Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
-    return getDocumentsMatchingQuery(query, offset, mutatedKeys, null);
+    return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
   }
 
   Iterable<Document> getDocuments() {

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

@@ -61,29 +61,35 @@
  */
 public class QueryEngine {
   private static final String LOG_TAG = "QueryEngine";
-  private static final int MIN_COLLECTION_SIZE_TO_AUTO_CREATE_INDEX = 100;
+
+  private static final int INDEX_AUTO_CREATION_DEFAULT_MIN_COLLECTION_SIZE = 100;
+
+  /**
+   * This cost represents the evaluation result of (([index, docKey] + [docKey, docContent]) per
+   * document in the result set) / ([docKey, docContent] per documents in full collection scan)
+   * coming from experiment https://github.com/firebase/firebase-android-sdk/pull/5064.
+   */
+  private static final double DEFAULT_RELATIVE_INDEX_READ_COST_PER_DOCUMENT = 2;
 
   private LocalDocumentsView localDocumentsView;
   private IndexManager indexManager;
   private boolean initialized;
 
-  private boolean automaticIndexingEnabled = false;
+  private boolean indexAutoCreationEnabled = false;
 
   /** SDK only decides whether it should create index when collection size is larger than this. */
-  private int minCollectionSizeToAutoCreateIndex = MIN_COLLECTION_SIZE_TO_AUTO_CREATE_INDEX;
+  private int indexAutoCreationMinCollectionSize = INDEX_AUTO_CREATION_DEFAULT_MIN_COLLECTION_SIZE;
+
+  private double relativeIndexReadCostPerDocument = DEFAULT_RELATIVE_INDEX_READ_COST_PER_DOCUMENT;
 
   public void initialize(LocalDocumentsView localDocumentsView, IndexManager indexManager) {
     this.localDocumentsView = localDocumentsView;
     this.indexManager = indexManager;
     this.initialized = true;
   }
 
-  public void enableIndexAutoCreation() {
-    this.automaticIndexingEnabled = true;
-  }
-
-  public void disableIndexAutoCreation() {
-    this.automaticIndexingEnabled = false;
+  public void setIndexAutoCreationEnabled(boolean enabled) {
+    this.indexAutoCreationEnabled = enabled;
   }
 
   public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
@@ -104,7 +110,7 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
 
     QueryContext context = new QueryContext();
     result = executeFullCollectionScan(query, context);
-    if (result != null && automaticIndexingEnabled) {
+    if (result != null && indexAutoCreationEnabled) {
       createCacheIndexes(query, context, result.size());
     }
     return result;
@@ -116,33 +122,32 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
    */
   // TODO(csi): Auto experiment data.
   private void createCacheIndexes(Query query, QueryContext context, int resultSize) {
-    if (context.getDocumentReadCount() < minCollectionSizeToAutoCreateIndex) {
+    if (context.getDocumentReadCount() < indexAutoCreationMinCollectionSize) {
       Logger.debug(
           LOG_TAG,
           "SDK will only creates cache indexes for collection contains more than or equal to "
               + "%s documents.",
-          minCollectionSizeToAutoCreateIndex);
+          indexAutoCreationMinCollectionSize);
       return;
     }
 
-    String decisionStr = "";
-    // If evaluation is updated, please update tests in SQLiteLocalStoreTest.java
-    if (context.getDocumentReadCount() > 2 * resultSize) {
+    Logger.debug(
+        LOG_TAG,
+        "Query scans %s local documents and returns %s documents as results.",
+        context.getDocumentReadCount(),
+        resultSize);
+
+    if (context.getDocumentReadCount() > relativeIndexReadCostPerDocument * resultSize) {
       indexManager.createTargetIndices(query.toTarget());
+      Logger.debug(
+          LOG_TAG,
+          "The SDK decides to create cache indexes for this query, as using cache indexes "
+              + "may help improve performance.");
     } else {
-      decisionStr = " not";
-    }
-
-    if (Logger.isDebugEnabled()) {
       Logger.debug(
           LOG_TAG,
-          "Query ran locally using a full collection scan, walking through %s documents in total "
-              + "and returning %s documents. The SDK has decided%s to create cache indexes "
-              + "for this query, as using cache indexes may%s help improve performance.",
-          context.getDocumentReadCount(),
-          resultSize,
-          decisionStr,
-          decisionStr);
+          "The SDK decides not to create cache indexes for this query, as using cache indexes "
+              + "may not help improve performance.");
     }
   }
 
@@ -321,7 +326,12 @@ private ImmutableSortedMap<DocumentKey, Document> appendRemainingResults(
   }
 
   @VisibleForTesting
-  void setMinCollectionSizeToAutoCreateIndex(int newMin) {
-    minCollectionSizeToAutoCreateIndex = newMin;
+  void setIndexAutoCreationMinCollectionSize(int newMin) {
+    indexAutoCreationMinCollectionSize = newMin;
+  }
+
+  @VisibleForTesting
+  void setRelativeIndexReadCostPerDocument(double newCost) {
+    relativeIndexReadCostPerDocument = newCost;
   }
 }

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

@@ -223,7 +223,6 @@ private Map<DocumentKey, MutableDocument> getAll(
             row -> {
               processRowInBackground(backgroundQueue, results, row, filter);
               if (context != null) {
-                // Increases the counter by 1 for every document processed.
                 context.incrementDocumentReadCount();
               }
             });
@@ -236,7 +235,7 @@ private Map<DocumentKey, MutableDocument> getAll(
       IndexOffset offset,
       int count,
       @Nullable Function<MutableDocument, Boolean> filter) {
-    return getAll(collections, offset, count, filter, null);
+    return getAll(collections, offset, count, filter, /*context*/ null);
   }
 
   private void processRowInBackground(
@@ -266,7 +265,7 @@ private void processRowInBackground(
   @Override
   public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
       Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
-    return getDocumentsMatchingQuery(query, offset, mutatedKeys, null);
+    return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
   }
 
   @Override

firebase-firestore/src/test/java/com/google/firebase/firestore/local/CountingQueryEngine.java

@@ -88,18 +88,18 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
   }
 
   @Override
-  public void enableIndexAutoCreation() {
-    queryEngine.enableIndexAutoCreation();
+  public void setIndexAutoCreationEnabled(boolean enabled) {
+    queryEngine.setIndexAutoCreationEnabled(enabled);
   }
 
   @Override
-  public void disableIndexAutoCreation() {
-    queryEngine.disableIndexAutoCreation();
+  public void setIndexAutoCreationMinCollectionSize(int newMin) {
+    queryEngine.setIndexAutoCreationMinCollectionSize(newMin);
   }
 
   @Override
-  public void setMinCollectionSizeToAutoCreateIndex(int newMin) {
-    queryEngine.setMinCollectionSizeToAutoCreateIndex(newMin);
+  public void setRelativeIndexReadCostPerDocument(double newCost) {
+    queryEngine.setRelativeIndexReadCostPerDocument(newCost);
   }
 
   /**
@@ -186,7 +186,7 @@ public Map<DocumentKey, MutableDocument> getAll(
       @Override
       public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
           Query query, IndexOffset offset, @NonNull Set<DocumentKey> mutatedKeys) {
-        return getDocumentsMatchingQuery(query, offset, mutatedKeys, null);
+        return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
       }
 
       @Override

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java

@@ -217,23 +217,21 @@ protected void executeQuery(Query query) {
     lastQueryResult = localStore.executeQuery(query, /* usePreviousResults= */ true);
   }
 
-  protected void enableIndexAutoCreation() {
+  protected void setIndexAutoCreationEnabled(boolean enabled) {
     // Noted: there are two queryEngines here, the first one is extended by CountingQueryEngine,
     // which is set by localStore function; The second one a pointer inside CountingQueryEngine,
     // which is set by queryEngine function.
     // Only the second function takes effect in the tests. Adding first one here for compatibility.
-    localStore.enableIndexAutoCreation();
-    queryEngine.enableIndexAutoCreation();
+    localStore.setIndexAutoCreationEnabled(enabled);
+    queryEngine.setIndexAutoCreationEnabled(enabled);
   }
 
-  protected void disableIndexAutoCreation() {
-    // Please refer to the notes in `enableIndexAutoCreation()`
-    localStore.disableIndexAutoCreation();
-    queryEngine.disableIndexAutoCreation();
+  protected void setMinCollectionSizeToAutoCreateIndex(int newMin) {
+    queryEngine.setIndexAutoCreationMinCollectionSize(newMin);
   }
 
-  protected void setMinCollectionSizeToAutoCreateIndex(int newMin) {
-    queryEngine.setMinCollectionSizeToAutoCreateIndex(newMin);
+  protected void setRelativeIndexReadCostPerDocument(double newCost) {
+    queryEngine.setRelativeIndexReadCostPerDocument(newCost);
   }
 
   private void releaseTarget(int targetId) {