Use the Tasks API

Author: schmidt-sebastian

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

@@ -17,6 +17,8 @@
 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;
@@ -30,9 +32,7 @@
 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 java.util.concurrent.ExecutionException;
 import javax.annotation.Nullable;
 
 final class SQLiteRemoteDocumentCache implements RemoteDocumentCache {
@@ -122,10 +122,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>> pendingDeserializations = new ArrayList<>();
 
     db.query("SELECT path, contents FROM remote_documents WHERE path >= ? AND path < ?")
         .binding(prefixPath, prefixSuccessorPath)
@@ -145,35 +142,43 @@ 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();
-                  });
+              pendingDeserializations.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;
+                      }));
             });
 
+    ImmutableSortedMap<DocumentKey, Document> matchingDocuments =
+        ImmutableSortedMap.Builder.emptyMap(DocumentKey.comparator());
+
     try {
-      completedTasks.acquire(pendingTaskCount[0]);
+      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);
+        }
+      }
+    } catch (ExecutionException e) {
+      throw new RuntimeException(e);
     } 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());
-      }
-    }
     return matchingDocuments;
   }