use threadlocal instead of synchronization

Author: dconeybe

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

@@ -61,29 +61,26 @@ public Overlay getOverlay(DocumentKey key) {
   @Override
   public Map<DocumentKey, Overlay> getOverlays(SortedSet<DocumentKey> keys) {
     hardAssert(keys.comparator() == null, "getOverlays() requires natural order");
-    Map<DocumentKey, Overlay> result = new HashMap<>();
 
-    BackgroundQueue backgroundQueue = new BackgroundQueue();
+    BackgroundQueue<DocumentKey, Overlay> backgroundQueue = new BackgroundQueue<>();
     ResourcePath currentCollection = ResourcePath.EMPTY;
     List<Object> accumulatedDocumentIds = new ArrayList<>();
     for (DocumentKey key : keys) {
       if (!currentCollection.equals(key.getCollectionPath())) {
-        processSingleCollection(result, backgroundQueue, currentCollection, accumulatedDocumentIds);
+        processSingleCollection(backgroundQueue, currentCollection, accumulatedDocumentIds);
         currentCollection = key.getCollectionPath();
         accumulatedDocumentIds.clear();
       }
       accumulatedDocumentIds.add(key.getDocumentId());
     }
 
-    processSingleCollection(result, backgroundQueue, currentCollection, accumulatedDocumentIds);
-    backgroundQueue.drain();
-    return result;
+    processSingleCollection(backgroundQueue, currentCollection, accumulatedDocumentIds);
+    return backgroundQueue.drain();
   }
 
   /** Reads the overlays for the documents in a single collection. */
   private void processSingleCollection(
-      Map<DocumentKey, Overlay> result,
-      BackgroundQueue backgroundQueue,
+      BackgroundQueue<DocumentKey, Overlay> backgroundQueue,
       ResourcePath collectionPath,
       List<Object> documentIds) {
     if (documentIds.isEmpty()) {
@@ -101,7 +98,7 @@ private void processSingleCollection(
     while (longQuery.hasMoreSubqueries()) {
       longQuery
           .performNextSubquery()
-          .forEach(row -> processOverlaysInBackground(backgroundQueue, result, row));
+          .forEach(row -> processOverlaysInBackground(backgroundQueue, row));
     }
   }
 
@@ -138,26 +135,23 @@ public void removeOverlaysForBatchId(int batchId) {
 
   @Override
   public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId) {
-    Map<DocumentKey, Overlay> result = new HashMap<>();
-    BackgroundQueue backgroundQueue = new BackgroundQueue();
+    BackgroundQueue<DocumentKey, Overlay> backgroundQueue = new BackgroundQueue<>();
     db.query(
             "SELECT overlay_mutation, largest_batch_id FROM document_overlays "
                 + "WHERE uid = ? AND collection_path = ? AND largest_batch_id > ?")
         .binding(uid, EncodedPath.encode(collection), sinceBatchId)
-        .forEach(row -> processOverlaysInBackground(backgroundQueue, result, row));
-    backgroundQueue.drain();
-    return result;
+        .forEach(row -> processOverlaysInBackground(backgroundQueue, row));
+    return backgroundQueue.drain();
   }
 
   @Override
   public Map<DocumentKey, Overlay> getOverlays(
       String collectionGroup, int sinceBatchId, int count) {
-    Map<DocumentKey, Overlay> result = new HashMap<>();
     String[] lastCollectionPath = new String[1];
     String[] lastDocumentPath = new String[1];
     int[] lastLargestBatchId = new int[1];
 
-    BackgroundQueue backgroundQueue = new BackgroundQueue();
+    BackgroundQueue<DocumentKey, Overlay> backgroundQueue1 = new BackgroundQueue<>();
     db.query(
             "SELECT overlay_mutation, largest_batch_id, collection_path, document_id "
                 + " FROM document_overlays "
@@ -169,17 +163,19 @@ public Map<DocumentKey, Overlay> getOverlays(
               lastLargestBatchId[0] = row.getInt(1);
               lastCollectionPath[0] = row.getString(2);
               lastDocumentPath[0] = row.getString(3);
-              processOverlaysInBackground(backgroundQueue, result, row);
+              processOverlaysInBackground(backgroundQueue1, row);
             });
 
+    HashMap<DocumentKey, Overlay> results = backgroundQueue1.drain();
     if (lastCollectionPath[0] == null) {
-      return result;
+      return results;
     }
 
     // This function should not return partial batch overlays, even if the number of overlays in the
     // result set exceeds the given `count` argument. Since the `LIMIT` in the above query might
     // result in a partial batch, the following query appends any remaining overlays for the last
     // batch.
+    BackgroundQueue<DocumentKey, Overlay> backgroundQueue2 = new BackgroundQueue<>();
     db.query(
             "SELECT overlay_mutation, largest_batch_id FROM document_overlays "
                 + "WHERE uid = ? AND collection_group = ? "
@@ -192,22 +188,21 @@ public Map<DocumentKey, Overlay> getOverlays(
             lastCollectionPath[0],
             lastDocumentPath[0],
             lastLargestBatchId[0])
-        .forEach(row -> processOverlaysInBackground(backgroundQueue, result, row));
-    backgroundQueue.drain();
-    return result;
+        .forEach(row -> processOverlaysInBackground(backgroundQueue2, row));
+
+    backgroundQueue2.drainInto(results);
+    return results;
   }
 
   private void processOverlaysInBackground(
-      BackgroundQueue backgroundQueue, Map<DocumentKey, Overlay> results, Cursor row) {
+      BackgroundQueue<DocumentKey, Overlay> backgroundQueue, Cursor row) {
     byte[] rawMutation = row.getBlob(0);
     int largestBatchId = row.getInt(1);
 
     backgroundQueue.submit(
-        () -> {
+        results -> {
           Overlay overlay = decodeOverlay(rawMutation, largestBatchId);
-          synchronized (results) {
-            results.put(overlay.getKey(), overlay);
-          }
+          results.put(overlay.getKey(), overlay);
         });
   }
 

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

@@ -168,21 +168,19 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> documentKe
             bindVars,
             ") ORDER BY path");
 
-    BackgroundQueue backgroundQueue = new BackgroundQueue();
+    BackgroundQueue<DocumentKey, MutableDocument> backgroundQueue = new BackgroundQueue<>();
     while (longQuery.hasMoreSubqueries()) {
       longQuery
           .performNextSubquery()
-          .forEach(row -> processRowInBackground(backgroundQueue, results, row, /*filter*/ null));
+          .forEach(row -> processRowInBackground(backgroundQueue, row, /*filter*/ null));
     }
-    backgroundQueue.drain();
+
+    backgroundQueue.drainInto(results);
 
     // Backfill any rows with null "document_type" discovered by processRowInBackground().
     documentTypeBackfiller.backfill(db);
 
-    // Synchronize on `results` to avoid a data race with the background queue.
-    synchronized (results) {
-      return results;
-    }
+    return results;
   }
 
   @Override
@@ -264,26 +262,23 @@ private Map<DocumentKey, MutableDocument> getAll(
     }
     bindVars[i] = count;
 
-    BackgroundQueue backgroundQueue = new BackgroundQueue();
-    Map<DocumentKey, MutableDocument> results = new HashMap<>();
+    BackgroundQueue<DocumentKey, MutableDocument> backgroundQueue = new BackgroundQueue<>();
     db.query(sql.toString())
         .binding(bindVars)
         .forEach(
             row -> {
-              processRowInBackground(backgroundQueue, results, row, filter);
+              processRowInBackground(backgroundQueue, row, filter);
               if (context != null) {
                 context.incrementDocumentReadCount();
               }
             });
-    backgroundQueue.drain();
+
+    HashMap<DocumentKey, MutableDocument> results = backgroundQueue.drain();
 
     // Backfill any null "document_type" columns discovered by processRowInBackground().
     documentTypeBackfiller.backfill(db);
 
-    // Synchronize on `results` to avoid a data race with the background queue.
-    synchronized (results) {
-      return results;
-    }
+    return results;
   }
 
   private Map<DocumentKey, MutableDocument> getAll(
@@ -296,8 +291,7 @@ private Map<DocumentKey, MutableDocument> getAll(
   }
 
   private void processRowInBackground(
-      BackgroundQueue backgroundQueue,
-      Map<DocumentKey, MutableDocument> results,
+      BackgroundQueue<DocumentKey, MutableDocument> backgroundQueue,
       Cursor row,
       @Nullable Function<MutableDocument, Boolean> filter) {
     byte[] rawDocument = row.getBlob(0);
@@ -307,16 +301,14 @@ private void processRowInBackground(
     String path = row.getString(4);
 
     backgroundQueue.submit(
-        () -> {
+        results -> {
           MutableDocument document =
               decodeMaybeDocument(rawDocument, readTimeSeconds, readTimeNanos);
           if (documentTypeIsNull) {
             documentTypeBackfiller.enqueue(path, readTimeSeconds, readTimeNanos, document);
           }
           if (filter == null || filter.apply(document)) {
-            synchronized (results) {
-              results.put(document.getKey(), document);
-            }
+            results.put(document.getKey(), document);
           }
         });
   }

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

@@ -27,24 +27,29 @@ import kotlinx.coroutines.asExecutor
  * called from the same thread. The behavior of an instance is undefined if any of the methods are
  * called from multiple threads.
  */
-internal class BackgroundQueue {
+internal class BackgroundQueue<K, V> {
 
-  private var state: State = State.Submitting()
+  private var state: State<K, V> = State.Submitting()
+
+  /** Overload for convenience of being called from Java code. */
+  fun submit(consumer: Consumer<HashMap<K, V>>) = this.submit { consumer.accept(it) }
 
   /**
    * Submit a task for asynchronous execution on the executor of the owning [BackgroundQueue].
    *
    * @throws IllegalStateException if [drain] has been called.
    */
-  fun submit(runnable: Runnable) {
+  fun submit(runnable: (results: HashMap<K, V>) -> Unit) {
     val submittingState = this.state
-    check(submittingState is State.Submitting) { "submit() may not be called after drain()" }
+    check(submittingState is State.Submitting) {
+      "submit() may not be called after drain() or drainInto()"
+    }
 
     submittingState.run {
       taskCount++
       executor.execute {
         try {
-          runnable.run()
+          runnable(threadLocalResults.get()!!)
         } finally {
           completedTasks.release()
         }
@@ -55,22 +60,53 @@ internal class BackgroundQueue {
   /**
    * Blocks until all tasks submitted via calls to [submit] have completed.
    *
-   * @throws IllegalStateException if called more than once.
+   * The results produced by each thread are merged into a new [HashMap] and returned.
+   *
+   * @throws IllegalStateException if [drain] or [drainInto] has already been called.
+   */
+  fun drain(): HashMap<K, V> = HashMap<K, V>().also { drainInto(it) }
+
+  /**
+   * Blocks until all tasks submitted via calls to [submit] have completed.
+   *
+   * The results produced by each thread are merged into the given map.
+   *
+   * @throws IllegalStateException if [drain] or [drainInto] has already been called.
    */
-  fun drain() {
+  fun drainInto(results: MutableMap<K, V>) {
     val submittingState = this.state
-    check(submittingState is State.Submitting) { "drain() may not be called more than once" }
-    this.state = State.Draining
+    check(submittingState is State.Submitting) { "drain() or drainInto() has already been called" }
+    this.state = State.Draining()
+    return submittingState.run {
+      completedTasks.acquire(taskCount)
+      threadLocalResults.mergeResultsInto(results)
+    }
+  }
+
+  private class ThreadLocalResults<K, V> : ThreadLocal<HashMap<K, V>>() {
+
+    private val results = mutableListOf<HashMap<K, V>>()
 
-    submittingState.run { completedTasks.acquire(taskCount) }
+    override fun initialValue(): HashMap<K, V> {
+      synchronized(results) {
+        val result = HashMap<K, V>()
+        results.add(result)
+        return result
+      }
+    }
+
+    fun mergeResultsInto(mergedResults: MutableMap<K, V>) {
+      synchronized(results) { results.forEach { mergedResults.putAll(it) } }
+    }
   }
 
-  private sealed interface State {
-    class Submitting : State {
+  private sealed interface State<K, V> {
+    class Submitting<K, V> : State<K, V> {
       val completedTasks = Semaphore(0)
+      val threadLocalResults = ThreadLocalResults<K, V>()
       var taskCount: Int = 0
     }
-    object Draining : State
+    class Draining<K, V> : State<K, V>
   }
 
   companion object {

firebase-firestore/src/test/java/com/google/firebase/firestore/core/QueryTest.java

@@ -45,7 +45,6 @@
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
-import java.util.Map;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -984,25 +983,19 @@ public void testSynchronousMatchesOrderBy() {
     // SQLiteRemoteDocumentCache.getAll(...), where `query.matches(doc)` is performed
     // for many different docs concurrently on the BackgroundQueue.
     Iterator<MutableDocument> iterator = docs.iterator();
-    BackgroundQueue backgroundQueue = new BackgroundQueue();
-    Map<DocumentKey, Boolean> results = new HashMap<>();
+    BackgroundQueue<DocumentKey, Boolean> backgroundQueue = new BackgroundQueue<>();
 
     while (iterator.hasNext()) {
       MutableDocument doc = iterator.next();
       backgroundQueue.submit(
-          () -> {
+          results -> {
             // We call query.matches() to indirectly test query.matchesOrderBy()
             boolean result = query.matches(doc);
-
-            // We will include a synchronized block in our command to simulate
-            // the implementation in SQLiteRemoteDocumentCache.getAll(...)
-            synchronized (results) {
-              results.put(doc.getKey(), result);
-            }
+            results.put(doc.getKey(), result);
           });
     }
 
-    backgroundQueue.drain();
+    HashMap<DocumentKey, Boolean> results = backgroundQueue.drain();
 
     Assert.assertEquals(101, results.keySet().size());
     for (DocumentKey key : results.keySet()) {