Make Task API work with tests

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 execution time 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

@@ -22,17 +22,18 @@
 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.ExecutionException;
 import javax.annotation.Nullable;
 
 final class SQLiteRemoteDocumentCache implements RemoteDocumentCache {
@@ -115,14 +116,17 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
         !query.isCollectionGroupQuery(),
         "CollectionGroup queries should be handled in LocalDocumentsView");
 
+    ImmutableSortedMap<DocumentKey, Document> matchingDocuments =
+        DocumentCollections.emptyDocumentMap();
+
     // Use the query path as a prefix for testing if a document matches the query.
     ResourcePath prefix = query.getPath();
     int immediateChildrenPathLength = prefix.length() + 1;
 
     String prefixPath = EncodedPath.encode(prefix);
     String prefixSuccessorPath = EncodedPath.prefixSuccessor(prefixPath);
 
-    List<Task<Document>> pendingDeserializations = new ArrayList<>();
+    List<Task<Document>> pendingTasks = new ArrayList<>();
 
     db.query("SELECT path, contents FROM remote_documents WHERE path >= ? AND path < ?")
         .binding(prefixPath, prefixSuccessorPath)
@@ -142,43 +146,29 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
 
               byte[] rawContents = row.getBlob(1);
 
-              pendingDeserializations.add(
+              pendingTasks.add(
                   Tasks.call(
                       // Since scheduling background tasks incurs overhead, we only dispatch to a
-                      // background
-                      // thread if there are still some documents remaining.
+                      // 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;
                       }));
             });
 
-    ImmutableSortedMap<DocumentKey, Document> matchingDocuments =
-        ImmutableSortedMap.Builder.emptyMap(DocumentKey.comparator());
-
-    try {
-      List<Task<Document>> result = Tasks.await(Tasks.whenAllSuccess(pendingDeserializations));
-      for (Task<Document> task : result) {
-        @Nullable Document doc = task.getResult();
-        if (doc != null) {
-          matchingDocuments = matchingDocuments.insert(doc.getKey(), doc);
-        }
+    List<Document> result = Util.await(Util.whenAllComplete(pendingTasks));
+    for (Document doc : result) {
+      if (doc != null) {
+        matchingDocuments = matchingDocuments.insert(doc.getKey(), doc);
       }
-    } catch (ExecutionException e) {
-      throw new RuntimeException(e);
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
     }
-
     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,7 +18,10 @@
 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.common.base.Function;
 import com.google.firebase.firestore.FieldPath;
 import com.google.firebase.firestore.FirebaseFirestoreException;
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
@@ -31,6 +34,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 +226,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);
+            });
+  }
 }