Use Task API

Author: schmidt-sebastian

firebase-firestore/CHANGELOG.md

@@ -1,5 +1,5 @@
 # Unreleased
-- [changed] Reduced execution time of queries with large result sets by up to
+- [changed] Improved performance of queries with large result sets by up to
   60%.
 - [changed] Instead of failing silently, Firestore now crashes the client app
   if it fails to load SSL Ciphers. To avoid these crashes, you must bundle 

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

@@ -35,11 +35,7 @@
 import java.util.List;
 import java.util.Map;
 
-/**
- * Serializer for values stored in the LocalStore.
- *
- * <p>This class is thread-safe.
- */
+/** Serializer for values stored in the LocalStore. */
 public final class LocalSerializer {
 
   private final RemoteSerializer rpcSerializer;

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

@@ -17,22 +17,23 @@
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
+import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.model.DocumentCollections;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MaybeDocument;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.util.Executors;
+import com.google.firebase.firestore.util.Util;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.Executor;
-import java.util.concurrent.Semaphore;
 import javax.annotation.Nullable;
 
 final class SQLiteRemoteDocumentCache implements RemoteDocumentCache {
@@ -122,10 +123,7 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
     String prefixPath = EncodedPath.encode(prefix);
     String prefixSuccessorPath = EncodedPath.prefixSuccessor(prefixPath);
 
-    Map<DocumentKey, Document> allDocuments = new ConcurrentHashMap<>();
-
-    int[] pendingTaskCount = new int[] {0};
-    Semaphore completedTasks = new Semaphore(0);
+    List<Task<Document>> pendingTasks = new ArrayList<>();
 
     db.query("SELECT path, contents FROM remote_documents WHERE path >= ? AND path < ?")
         .binding(prefixPath, prefixSuccessorPath)
@@ -145,33 +143,29 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
 
               byte[] rawContents = row.getBlob(1);
 
-              ++pendingTaskCount[0];
-
-              // Since scheduling background tasks incurs overhead, we only dispatch to a background
-              // thread if there are still some documents remaining.
-              Executor deserializationExecutor =
-                  row.isLast() ? Executors.DIRECT_EXECUTOR : Executors.BACKGROUND_EXECUTOR;
-              deserializationExecutor.execute(
-                  () -> {
-                    MaybeDocument maybeDoc = decodeMaybeDocument(rawContents);
-                    if (maybeDoc instanceof Document) {
-                      allDocuments.put(maybeDoc.getKey(), (Document) maybeDoc);
-                    }
-                    completedTasks.release();
-                  });
+              pendingTasks.add(
+                  Tasks.call(
+                      // Since scheduling background tasks incurs overhead, we only dispatch to a
+                      // background thread if there are still some documents remaining.
+                      row.isLast() ? Executors.DIRECT_EXECUTOR : Executors.BACKGROUND_EXECUTOR,
+                      () -> {
+                        MaybeDocument maybeDoc = decodeMaybeDocument(rawContents);
+                        if (!(maybeDoc instanceof Document)) {
+                          return null;
+                        }
+                        if (!query.matches((Document) maybeDoc)) {
+                          return null;
+                        }
+                        return (Document) maybeDoc;
+                      }));
             });
 
-    try {
-      completedTasks.acquire(pendingTaskCount[0]);
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-    }
-
     ImmutableSortedMap<DocumentKey, Document> matchingDocuments =
-        ImmutableSortedMap.Builder.emptyMap(DocumentKey.comparator());
-    for (Map.Entry<DocumentKey, Document> entry : allDocuments.entrySet()) {
-      if (query.matches(entry.getValue())) {
-        matchingDocuments = matchingDocuments.insert(entry.getKey(), entry.getValue());
+        DocumentCollections.emptyDocumentMap();
+    List<Document> results = Util.await(Util.whenAllComplete(pendingTasks));
+    for (Document doc : results) {
+      if (doc != null) {
+        matchingDocuments = matchingDocuments.insert(doc.getKey(), doc);
       }
     }
     return matchingDocuments;

firebase-firestore/src/main/java/com/google/firebase/firestore/model/DatabaseId.java

@@ -16,11 +16,7 @@
 
 import androidx.annotation.NonNull;
 
-/**
- * Represents a particular database in Firestore.
- *
- * <p>This class is thread-safe.
- */
+/** Represents a particular database in Firestore */
 public final class DatabaseId implements Comparable<DatabaseId> {
   public static final String DEFAULT_DATABASE_ID = "(default)";
 

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteSerializer.java

@@ -106,11 +106,7 @@
 import java.util.Map;
 import java.util.Set;
 
-/**
- * Serializer that converts to and from Firestore API protos.
- *
- * <p>This class is thread-safe.
- */
+/** Serializer that converts to and from Firestore API protos. */
 public final class RemoteSerializer {
 
   private final DatabaseId databaseId;

firebase-firestore/src/main/java/com/google/firebase/firestore/util/Util.java

@@ -18,6 +18,8 @@
 import android.os.Looper;
 import androidx.annotation.Nullable;
 import com.google.android.gms.tasks.Continuation;
+import com.google.android.gms.tasks.Task;
+import com.google.android.gms.tasks.Tasks;
 import com.google.cloud.datastore.core.number.NumberComparisonHelper;
 import com.google.firebase.firestore.FieldPath;
 import com.google.firebase.firestore.FirebaseFirestoreException;
@@ -31,6 +33,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Random;
+import java.util.concurrent.Semaphore;
 
 /** A utility class for Firestore */
 public class Util {
@@ -222,4 +225,39 @@ public static void crashMainThread(RuntimeException exception) {
               throw exception;
             });
   }
+
+  /**
+   * Waits for the given Task to complete.
+   *
+   * <p>Similar to Tasks.await() but can be called from the main thread to support unit testing.
+   */
+  public static <TResult> TResult await(Task<TResult> task) {
+    Semaphore semaphore = new Semaphore(0);
+    task.addOnCompleteListener(Executors.BACKGROUND_EXECUTOR, t -> semaphore.release());
+    try {
+      semaphore.acquire();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+    return task.getResult();
+  }
+
+  /**
+   * Returns a Task with a list of Tasks that completes successfully when all of the specified Tasks
+   * complete.
+   *
+   * <p>Similar to Tasks.whenAllComplete() but does not schedule the completion on the main thread.
+   */
+  public static <TResult> Task<List<TResult>> whenAllComplete(List<Task<TResult>> tasks) {
+    return Tasks.whenAll(tasks)
+        .continueWithTask(
+            Executors.BACKGROUND_EXECUTOR,
+            t -> {
+              List<TResult> results = new ArrayList<>();
+              for (Task<TResult> task : tasks) {
+                results.add(task.getResult());
+              }
+              return Tasks.forResult(results);
+            });
+  }
 }