Revert "use threadlocal instead of synchronization"

This reverts commit 32c3add.

It made no performance improvements whatsoever.

Author: dconeybe

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

@@ -61,26 +61,29 @@ 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<DocumentKey, Overlay> backgroundQueue = new BackgroundQueue<>();
+    BackgroundQueue backgroundQueue = new BackgroundQueue();
     ResourcePath currentCollection = ResourcePath.EMPTY;
     List<Object> accumulatedDocumentIds = new ArrayList<>();
     for (DocumentKey key : keys) {
       if (!currentCollection.equals(key.getCollectionPath())) {
-        processSingleCollection(backgroundQueue, currentCollection, accumulatedDocumentIds);
+        processSingleCollection(result, backgroundQueue, currentCollection, accumulatedDocumentIds);
         currentCollection = key.getCollectionPath();
         accumulatedDocumentIds.clear();
       }
       accumulatedDocumentIds.add(key.getDocumentId());
     }
 
-    processSingleCollection(backgroundQueue, currentCollection, accumulatedDocumentIds);
-    return backgroundQueue.drain();
+    processSingleCollection(result, backgroundQueue, currentCollection, accumulatedDocumentIds);
+    backgroundQueue.drain();
+    return result;
   }
 
   /** Reads the overlays for the documents in a single collection. */
   private void processSingleCollection(
-      BackgroundQueue<DocumentKey, Overlay> backgroundQueue,
+      Map<DocumentKey, Overlay> result,
+      BackgroundQueue backgroundQueue,
       ResourcePath collectionPath,
       List<Object> documentIds) {
     if (documentIds.isEmpty()) {
@@ -98,7 +101,7 @@ private void processSingleCollection(
     while (longQuery.hasMoreSubqueries()) {
       longQuery
           .performNextSubquery()
-          .forEach(row -> processOverlaysInBackground(backgroundQueue, row));
+          .forEach(row -> processOverlaysInBackground(backgroundQueue, result, row));
     }
   }
 
@@ -135,23 +138,26 @@ public void removeOverlaysForBatchId(int batchId) {
 
   @Override
   public Map<DocumentKey, Overlay> getOverlays(ResourcePath collection, int sinceBatchId) {
-    BackgroundQueue<DocumentKey, Overlay> backgroundQueue = new BackgroundQueue<>();
+    Map<DocumentKey, Overlay> result = new HashMap<>();
+    BackgroundQueue 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, row));
-    return backgroundQueue.drain();
+        .forEach(row -> processOverlaysInBackground(backgroundQueue, result, row));
+    backgroundQueue.drain();
+    return result;
   }
 
   @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<DocumentKey, Overlay> backgroundQueue1 = new BackgroundQueue<>();
+    BackgroundQueue backgroundQueue = new BackgroundQueue();
     db.query(
             "SELECT overlay_mutation, largest_batch_id, collection_path, document_id "
                 + " FROM document_overlays "
@@ -163,19 +169,17 @@ public Map<DocumentKey, Overlay> getOverlays(
               lastLargestBatchId[0] = row.getInt(1);
               lastCollectionPath[0] = row.getString(2);
               lastDocumentPath[0] = row.getString(3);
-              processOverlaysInBackground(backgroundQueue1, row);
+              processOverlaysInBackground(backgroundQueue, result, row);
             });
 
-    HashMap<DocumentKey, Overlay> results = backgroundQueue1.drain();
     if (lastCollectionPath[0] == null) {
-      return results;
+      return result;
     }
 
     // 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 = ? "
@@ -188,21 +192,22 @@ public Map<DocumentKey, Overlay> getOverlays(
             lastCollectionPath[0],
             lastDocumentPath[0],
             lastLargestBatchId[0])
-        .forEach(row -> processOverlaysInBackground(backgroundQueue2, row));
-
-    backgroundQueue2.drainInto(results);
-    return results;
+        .forEach(row -> processOverlaysInBackground(backgroundQueue, result, row));
+    backgroundQueue.drain();
+    return result;
   }
 
   private void processOverlaysInBackground(
-      BackgroundQueue<DocumentKey, Overlay> backgroundQueue, Cursor row) {
+      BackgroundQueue backgroundQueue, Map<DocumentKey, Overlay> results, Cursor row) {
     byte[] rawMutation = row.getBlob(0);
     int largestBatchId = row.getInt(1);
 
     backgroundQueue.submit(
-        results -> {
+        () -> {
           Overlay overlay = decodeOverlay(rawMutation, largestBatchId);
-          results.put(overlay.getKey(), overlay);
+          synchronized (results) {
+            results.put(overlay.getKey(), overlay);
+          }
         });
   }
 

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

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

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

@@ -27,29 +27,24 @@ 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<K, V> {
+internal class BackgroundQueue {
 
-  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) }
+  private var state: State = State.Submitting()
 
   /**
    * Submit a task for asynchronous execution on the executor of the owning [BackgroundQueue].
    *
    * @throws IllegalStateException if [drain] has been called.
    */
-  fun submit(runnable: (results: HashMap<K, V>) -> Unit) {
+  fun submit(runnable: Runnable) {
     val submittingState = this.state
-    check(submittingState is State.Submitting) {
-      "submit() may not be called after drain() or drainInto()"
-    }
+    check(submittingState is State.Submitting) { "submit() may not be called after drain()" }
 
     submittingState.run {
       taskCount++
       executor.execute {
         try {
-          runnable(threadLocalResults.get()!!)
+          runnable.run()
         } finally {
           completedTasks.release()
         }
@@ -60,53 +55,22 @@ internal class BackgroundQueue<K, V> {
   /**
    * Blocks until all tasks submitted via calls to [submit] have completed.
    *
-   * 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.
+   * @throws IllegalStateException if called more than once.
    */
-  fun drainInto(results: MutableMap<K, V>) {
+  fun drain() {
     val submittingState = this.state
-    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>>()
+    check(submittingState is State.Submitting) { "drain() may not be called more than once" }
+    this.state = State.Draining
 
-    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) } }
-    }
+    submittingState.run { completedTasks.acquire(taskCount) }
   }
 
-  private sealed interface State<K, V> {
-    class Submitting<K, V> : State<K, V> {
+  private sealed interface State {
+    class Submitting : State {
       val completedTasks = Semaphore(0)
-      val threadLocalResults = ThreadLocalResults<K, V>()
       var taskCount: Int = 0
     }
-    class Draining<K, V> : State<K, V>
+    object Draining : State
   }
 
   companion object {

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

@@ -45,6 +45,7 @@
 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;
@@ -983,19 +984,25 @@ 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<DocumentKey, Boolean> backgroundQueue = new BackgroundQueue<>();
+    BackgroundQueue backgroundQueue = new BackgroundQueue();
+    Map<DocumentKey, Boolean> results = new HashMap<>();
 
     while (iterator.hasNext()) {
       MutableDocument doc = iterator.next();
       backgroundQueue.submit(
-          results -> {
+          () -> {
             // We call query.matches() to indirectly test query.matchesOrderBy()
             boolean result = query.matches(doc);
-            results.put(doc.getKey(), result);
+
+            // We will include a synchronized block in our command to simulate
+            // the implementation in SQLiteRemoteDocumentCache.getAll(...)
+            synchronized (results) {
+              results.put(doc.getKey(), result);
+            }
           });
     }
 
-    HashMap<DocumentKey, Boolean> results = backgroundQueue.drain();
+    backgroundQueue.drain();
 
     Assert.assertEquals(101, results.keySet().size());
     for (DocumentKey key : results.keySet()) {