Performance: Decode documents in background thread (#559)

Author: schmidt-sebastian

firebase-firestore/CHANGELOG.md

@@ -1,7 +1,8 @@
 # Unreleased
-- [fixed] Fixed an internal assertion that was triggered when an
-  update with a `FieldValue.serverTimestamp()` and an update with a
+- [fixed] Fixed an internal assertion that was triggered when an update
+  with a `FieldValue.serverTimestamp()` and an update with a
   `FieldValue.increment()` were pending for the same document (#491).
+- [changed] Improved performance of queries with large result sets.
 - [changed] Improved performance for queries with filters that only return a
   small subset of the documents in a collection.
 - [changed] Instead of failing silently, Firestore now crashes the client app

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

@@ -20,15 +20,19 @@
 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.BackgroundQueue;
+import com.google.firebase.firestore.util.Executors;
 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.Executor;
 import javax.annotation.Nullable;
 
 final class SQLiteRemoteDocumentCache implements RemoteDocumentCache {
@@ -118,7 +122,11 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
     String prefixPath = EncodedPath.encode(prefix);
     String prefixSuccessorPath = EncodedPath.prefixSuccessor(prefixPath);
 
-    Map<DocumentKey, Document> results = new HashMap<>();
+    BackgroundQueue backgroundQueue = new BackgroundQueue();
+
+    ImmutableSortedMap<DocumentKey, Document>[] matchingDocuments =
+        (ImmutableSortedMap<DocumentKey, Document>[])
+            new ImmutableSortedMap[] {DocumentCollections.emptyDocumentMap()};
 
     db.query("SELECT path, contents FROM remote_documents WHERE path >= ? AND path < ?")
         .binding(prefixPath, prefixSuccessorPath)
@@ -136,20 +144,31 @@ public ImmutableSortedMap<DocumentKey, Document> getAllDocumentsMatchingQuery(Qu
                 return;
               }
 
-              MaybeDocument maybeDoc = decodeMaybeDocument(row.getBlob(1));
-              if (!(maybeDoc instanceof Document)) {
-                return;
-              }
-
-              Document doc = (Document) maybeDoc;
-              if (!query.matches(doc)) {
-                return;
-              }
-
-              results.put(doc.getKey(), doc);
+              byte[] rawDocument = row.getBlob(1);
+
+              // Since scheduling background tasks incurs overhead, we only dispatch to a background
+              // thread if there are still some documents remaining.
+              Executor executor = row.isLast() ? Executors.DIRECT_EXECUTOR : backgroundQueue;
+              executor.execute(
+                  () -> {
+                    MaybeDocument maybeDoc = decodeMaybeDocument(rawDocument);
+
+                    if (maybeDoc instanceof Document && query.matches((Document) maybeDoc)) {
+                      synchronized (SQLiteRemoteDocumentCache.this) {
+                        matchingDocuments[0] =
+                            matchingDocuments[0].insert(maybeDoc.getKey(), (Document) maybeDoc);
+                      }
+                    }
+                  });
             });
 
-    return ImmutableSortedMap.Builder.fromMap(results, DocumentKey.comparator());
+    try {
+      backgroundQueue.drain();
+    } catch (InterruptedException e) {
+      fail("Interrupted while deserializing documents", e);
+    }
+
+    return matchingDocuments[0];
   }
 
   private String pathForKey(DocumentKey key) {

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

@@ -15,7 +15,6 @@
 package com.google.firebase.firestore.remote;
 
 import android.content.Context;
-import android.os.AsyncTask;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.common.GooglePlayServicesNotAvailableException;
 import com.google.android.gms.common.GooglePlayServicesRepairableException;
@@ -24,6 +23,7 @@
 import com.google.android.gms.tasks.Tasks;
 import com.google.firebase.firestore.core.DatabaseInfo;
 import com.google.firebase.firestore.util.AsyncQueue;
+import com.google.firebase.firestore.util.Executors;
 import com.google.firebase.firestore.util.Logger;
 import com.google.firebase.firestore.util.Supplier;
 import com.google.firestore.v1.FirestoreGrpc;
@@ -69,11 +69,11 @@ public static void overrideChannelBuilder(
       CallCredentials firestoreHeaders) {
     this.asyncQueue = asyncQueue;
 
-    // We execute network initialization on a separate thred to not block operations that depend on
+    // We execute network initialization on a separate thread to not block operations that depend on
     // the AsyncQueue.
     this.channelTask =
         Tasks.call(
-            AsyncTask.THREAD_POOL_EXECUTOR,
+            Executors.BACKGROUND_EXECUTOR,
             () -> {
               ManagedChannel channel = initChannel(context, databaseInfo);
               FirestoreGrpc.FirestoreStub firestoreStub =

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

@@ -0,0 +1,47 @@
+// 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 java.util.concurrent.Executor;
+import java.util.concurrent.Semaphore;
+
+/**
+ * A simple queue that executes tasks in parallel on the Android's AsyncTask.THREAD_POOL_EXECUTOR
+ * and supports blocking on their completion.
+ *
+ * <p>This class is not thread-safe. In particular, `execute()` and `drain()` should not be called
+ * from parallel threads.
+ */
+public class BackgroundQueue implements Executor {
+  private Semaphore completedTasks = new Semaphore(0);
+  private int pendingTaskCount = 0;
+
+  /** Enqueue a task on Android's THREAD_POOL_EXECUTOR. */
+  @Override
+  public void execute(Runnable task) {
+    ++pendingTaskCount;
+    Executors.BACKGROUND_EXECUTOR.execute(
+        () -> {
+          task.run();
+          completedTasks.release();
+        });
+  }
+
+  /** Wait for all currently scheduled tasks to complete. */
+  public void drain() throws InterruptedException {
+    completedTasks.acquire(pendingTaskCount);
+    pendingTaskCount = 0;
+  }
+}

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

@@ -14,11 +14,20 @@
 
 package com.google.firebase.firestore.util;
 
+import android.os.AsyncTask;
 import com.google.android.gms.tasks.TaskExecutors;
 import java.util.concurrent.Executor;
 
 /** Helper class for executors. */
 public final class Executors {
+  /**
+   * The maximum number of tasks we submit to AsyncTask.THREAD_POOL_EXECUTOR.
+   *
+   * <p>The limit is based on the number of core threads spun by THREAD_POOL_EXECUTOR and is well
+   * below the queue size limit of 120 pending tasks. Limiting our usage of the THREAD_POOL_EXECUTOR
+   * allows other users to schedule their own operations on the shared THREAD_POOL_EXECUTOR.
+   */
+  private static final int ASYNC_THREAD_POOL_MAXIMUM_CONCURRENCY = 4;
 
   /**
    * The default executor for user visible callbacks. It is an executor scheduling callbacks on
@@ -29,6 +38,11 @@ public final class Executors {
   /** An executor that executes the provided runnable immediately on the current thread. */
   public static final Executor DIRECT_EXECUTOR = Runnable::run;
 
+  /** An executor that runs tasks in parallel on Android's AsyncTask.THREAD_POOL_EXECUTOR. */
+  public static final Executor BACKGROUND_EXECUTOR =
+      new ThrottledForwardingExecutor(
+          ASYNC_THREAD_POOL_MAXIMUM_CONCURRENCY, AsyncTask.THREAD_POOL_EXECUTOR);
+
   private Executors() {
     // Private constructor to prevent initialization
   }

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

@@ -0,0 +1,64 @@
+// 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 java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.Semaphore;
+
+/**
+ * An executor that forwards executions to another executor, but caps the number of pending
+ * operations. Tasks scheduled past the specified limit are directly invoked on the calling thread,
+ * reducing the total memory consumed by pending tasks.
+ */
+class ThrottledForwardingExecutor implements Executor {
+  private final Executor executor;
+  private final Semaphore availableSlots;
+
+  /**
+   * Instantiates a new ThrottledForwardingExecutor.
+   *
+   * @param maximumConcurrency The maximum number of pending tasks to schedule on the provided
+   *     executor.
+   * @param executor The executor to forward tasks to.
+   */
+  ThrottledForwardingExecutor(int maximumConcurrency, Executor executor) {
+    this.availableSlots = new Semaphore(maximumConcurrency);
+    this.executor = executor;
+  }
+
+  /**
+   * Forwards a task to the provided executor if the current number of pending tasks is less than
+   * the configured limit. Otherwise, executes the task directly.
+   *
+   * @param command The task to run.
+   */
+  @Override
+  public void execute(Runnable command) {
+    if (availableSlots.tryAcquire()) {
+      try {
+        executor.execute(
+            () -> {
+              command.run();
+              availableSlots.release();
+            });
+      } catch (RejectedExecutionException e) {
+        command.run();
+      }
+    } else {
+      command.run();
+    }
+  }
+}

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

@@ -0,0 +1,99 @@
+// 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 static org.junit.Assert.fail;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.junit.Test;
+
+public class ThrottledForwardingExecutorTest {
+  @Test
+  public void limitsNumberOfForwardedTasks() throws InterruptedException {
+    Semaphore completedTasks = new Semaphore(0);
+    int maximumConcurrency = 10;
+
+    CountingExecutor countingExecutor = new CountingExecutor();
+    ThrottledForwardingExecutor throttledExecutor =
+        new ThrottledForwardingExecutor(maximumConcurrency, countingExecutor);
+
+    // Schedule more than `maximumConcurrency` parallel tasks and wait until all scheduling has
+    // finished.
+    int numTasks = maximumConcurrency + 1;
+    CountDownLatch schedulingLatch = new CountDownLatch(1);
+    for (int i = 0; i < numTasks; ++i) {
+      int currentTask = i;
+      throttledExecutor.execute(
+          () -> {
+            try {
+              if (currentTask < maximumConcurrency) {
+                // Block if we are running on the forwarded executor. We can't block the thread that
+                // is running this test.
+                schedulingLatch.await();
+              }
+              completedTasks.release();
+            } catch (InterruptedException e) {
+              fail("Unexpected InterruptedException: " + e);
+            }
+          });
+    }
+    schedulingLatch.countDown();
+
+    // Verify that only `maximumConcurrency` tasks were forwarded to the executor.
+    completedTasks.acquire(numTasks);
+    assertEquals(maximumConcurrency, countingExecutor.getNumTasks());
+  }
+
+  @Test
+  public void handlesRejectedExecutionException() {
+    AtomicInteger result = new AtomicInteger(0);
+
+    ThrottledForwardingExecutor executor =
+        new ThrottledForwardingExecutor(
+            10,
+            command -> {
+              throw new RejectedExecutionException();
+            });
+
+    executor.execute(result::incrementAndGet);
+
+    assertEquals(1, result.get());
+  }
+
+  /** An executor that counts the number of tasks submitted. */
+  private static class CountingExecutor implements Executor {
+    int numTasks = 0;
+
+    @Override
+    public void execute(Runnable command) {
+      ++numTasks;
+      new Thread() {
+        @Override
+        public void run() {
+          command.run();
+        }
+      }.start();
+    }
+
+    int getNumTasks() {
+      return numTasks;
+    }
+  }
+}