Realize Android Tasks are slow

Author: schmidt-sebastian

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

@@ -17,23 +17,21 @@
 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.firebase.firestore.util.TaskQueue;
 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 {
@@ -123,7 +121,7 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
     String prefixPath = EncodedPath.encode(prefix);
     String prefixSuccessorPath = EncodedPath.prefixSuccessor(prefixPath);
 
-    List<Task<Document>> pendingTasks = new ArrayList<>();
+    TaskQueue<Document> taskQueue = new TaskQueue<>();
 
     db.query("SELECT path, contents FROM remote_documents WHERE path >= ? AND path < ?")
         .binding(prefixPath, prefixSuccessorPath)
@@ -141,40 +139,48 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
                 return;
               }
 
-              byte[] rawContents = row.getBlob(1);
-
-              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;
-                      }));
+              byte[] rawDocument = row.getBlob(1);
+
+              if (!row.isLast()) {
+                taskQueue.enqueueInBackground(() -> this.processDocumentResult(rawDocument, query));
+              } else {
+                // Enqueue the last document inline so that we don't block on background execution.
+                taskQueue.enqueueInline(() -> this.processDocumentResult(rawDocument, query));
+              }
             });
 
     ImmutableSortedMap<DocumentKey, Document> matchingDocuments =
         DocumentCollections.emptyDocumentMap();
-    List<Document> results = Util.await(Util.whenAllComplete(pendingTasks));
-    for (Document doc : results) {
-      if (doc != null) {
-        matchingDocuments = matchingDocuments.insert(doc.getKey(), doc);
+
+    try {
+      List<Document> results = taskQueue.awaitResults();
+      for (Document doc : results) {
+        if (doc != null) {
+          matchingDocuments = matchingDocuments.insert(doc.getKey(), doc);
+        }
       }
+    } catch (ExecutionException e) {
+      fail(e, "Failed to decode documents from SQLite");
     }
+
     return matchingDocuments;
   }
 
   private String pathForKey(DocumentKey key) {
     return EncodedPath.encode(key.getPath());
   }
 
+  private @Nullable Document processDocumentResult(byte[] rawDocument, Query query) {
+    MaybeDocument maybeDoc = decodeMaybeDocument(rawDocument);
+    if (!(maybeDoc instanceof Document)) {
+      return null;
+    }
+    if (!query.matches((Document) maybeDoc)) {
+      return null;
+    }
+    return (Document) maybeDoc;
+  }
+
   private MaybeDocument decodeMaybeDocument(byte[] bytes) {
     try {
       return serializer.decodeMaybeDocument(

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

@@ -0,0 +1,120 @@
+// Copyright 2019 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.util;
+
+import static com.google.firebase.firestore.util.Assert.hardAssert;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.Semaphore;
+
+/**
+ * A queue for parallel execution of independent tasks.
+ *
+ * <p>All tasks must return the same result type T and results can be be fetched via
+ * `awaitResults()`. It is an error to enqueue tasks after invoking `awaitResults()`.
+ *
+ * <p>This class is not thread-safe. All public methods must be called from the same thread.
+ */
+public class TaskQueue<T> {
+  private BlockingQueue<TaskResult<T>> taskResults = new PriorityBlockingQueue<>();
+  private Semaphore completedTasks = new Semaphore(0);
+  private int totalTaskCount = 0;
+  private boolean acceptsTasks = true;
+
+  /** A single result from a background task. Encapsulates the task number to allow for sorting. */
+  private static class TaskResult<T> implements Comparable<TaskResult<T>> {
+    final int numTask;
+    final T result;
+    final Exception exception;
+
+    private TaskResult(int numTask, T result) {
+      this.numTask = numTask;
+      this.result = result;
+      this.exception = null;
+    }
+
+    private TaskResult(int numTask, Exception exception) {
+      this.numTask = numTask;
+      this.result = null;
+      this.exception = exception;
+    }
+
+    @Override
+    public int compareTo(TaskResult<T> other) {
+      return Integer.compare(numTask, other.numTask);
+    }
+  }
+
+  /** Enqueue a task on Android's THREAD_POOL_EXECUTOR. */
+  public void enqueueInBackground(Callable<T> task) {
+    hardAssert(acceptsTasks, "enqueueInBackground() called after awaitResults()");
+
+    int currentTask = ++totalTaskCount;
+    Executors.BACKGROUND_EXECUTOR.execute(
+        () -> {
+          try {
+            T result = task.call();
+            taskResults.add(new TaskResult<>(currentTask, result));
+          } catch (Exception e) {
+            taskResults.add(new TaskResult<>(currentTask, e));
+          }
+          completedTasks.release();
+        });
+  }
+
+  /** Executes a task inline, allowing for more efficient execution of small tasks. */
+  public void enqueueInline(Callable<T> callable) {
+    hardAssert(acceptsTasks, "enqueueInBackground() called after awaitResults()");
+
+    int currentTask = ++totalTaskCount;
+    try {
+      T result = callable.call();
+      taskResults.add(new TaskResult<>(currentTask, result));
+    } catch (Exception e) {
+      taskResults.add(new TaskResult<>(currentTask, e));
+    }
+    completedTasks.release();
+  }
+
+  /**
+   * Blocks on the execution of all tasks. Returns a list of results in the order that tasks where
+   * added.
+   */
+  public List<T> awaitResults() throws ExecutionException {
+    acceptsTasks = false;
+
+    List<T> allResults = new ArrayList<>(totalTaskCount);
+    try {
+      completedTasks.acquire(totalTaskCount);
+      while (!taskResults.isEmpty()) {
+        TaskResult<T> currentResult = taskResults.take();
+
+        if (currentResult.exception != null) {
+          throw new ExecutionException("Unhandled exception in task", currentResult.exception);
+        } else {
+          allResults.add(currentResult.result);
+        }
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+    return allResults;
+  }
+}

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

@@ -18,8 +18,6 @@
 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;
@@ -33,7 +31,6 @@
 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 {
@@ -225,39 +222,4 @@ 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);
-            });
-  }
 }

firebase-firestore/src/test/java/com/google/firebase/firestore/util/TaskQueueTest.java

@@ -0,0 +1,107 @@
+// Copyright 2019 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.util;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Semaphore;
+import org.junit.Test;
+import org.robolectric.annotation.Config;
+
+@org.junit.runner.RunWith(org.robolectric.RobolectricTestRunner.class)
+@Config(manifest = Config.NONE)
+public class TaskQueueTest {
+
+  @Test
+  public void canEnqueueInBackground() throws ExecutionException {
+    TaskQueue<Integer> queue = new TaskQueue<>();
+    queue.enqueueInBackground(() -> 1);
+    queue.enqueueInBackground(() -> 2);
+    queue.enqueueInBackground(() -> 3);
+    List<Integer> result = queue.awaitResults();
+    assertEquals(Arrays.asList(1, 2, 3), result);
+  }
+
+  @Test
+  public void canEnqueueInline() throws ExecutionException {
+    TaskQueue<Integer> queue = new TaskQueue<>();
+    queue.enqueueInline(() -> 1);
+    queue.enqueueInline(() -> 2);
+    queue.enqueueInline(() -> 3);
+    List<Integer> result = queue.awaitResults();
+    assertEquals(Arrays.asList(1, 2, 3), result);
+  }
+
+  @Test(expected = ExecutionException.class)
+  public void enqueueInBackgroundForwardsException() throws ExecutionException {
+    TaskQueue<Integer> queue = new TaskQueue<>();
+    queue.enqueueInBackground(() -> 1);
+    queue.enqueueInBackground(
+        () -> {
+          throw new Exception("foo");
+        });
+    queue.enqueueInBackground(() -> 3);
+    queue.awaitResults();
+  }
+
+  @Test(expected = ExecutionException.class)
+  public void enqueueInlineForwardsException() throws ExecutionException {
+    TaskQueue<Integer> queue = new TaskQueue<>();
+    queue.enqueueInline(() -> 1);
+    queue.enqueueInline(
+        () -> {
+          throw new Exception("foo");
+        });
+    queue.enqueueInline(() -> 3);
+    queue.awaitResults();
+  }
+
+  @Test
+  public void enqueuePreservesOrder() throws ExecutionException {
+    Semaphore firstTaskCompleted = new Semaphore(0);
+    Semaphore secondTaskCompleted = new Semaphore(0);
+    Semaphore thirdTaskCompleted = new Semaphore(0);
+
+    TaskQueue<Integer> queue = new TaskQueue<>();
+    queue.enqueueInBackground(
+        () -> {
+          thirdTaskCompleted.acquire();
+          return 1;
+        });
+    queue.enqueueInBackground(
+        () -> {
+          secondTaskCompleted.acquire();
+          thirdTaskCompleted.release();
+          return 2;
+        });
+    queue.enqueueInBackground(
+        () -> {
+          firstTaskCompleted.release();
+          return 3;
+        });
+    queue.enqueueInline(
+        () -> {
+          firstTaskCompleted.acquire();
+          secondTaskCompleted.release();
+          return 4;
+        });
+
+    List<Integer> result = queue.awaitResults();
+    assertEquals(Arrays.asList(1, 2, 3, 4), result);
+  }
+}