Implement limbo resolution listen throttling.

Author: dconeybe

firebase-firestore/src/main/java/com/google/firebase/firestore/core/SyncEngine.java

@@ -17,6 +17,7 @@
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.android.gms.tasks.Task;
@@ -48,11 +49,13 @@
 import com.google.firebase.firestore.util.Logger;
 import com.google.firebase.firestore.util.Util;
 import io.grpc.Status;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Queue;
 import java.util.Set;
 
 /**
@@ -89,8 +92,28 @@ private static class LimboResolution {
     }
   }
 
+  /** An entry in {@link #limboListenQueue}. */
+  private static final class LimboListenQueueEntry {
+
+    final DocumentKey key;
+    final TargetData targetData;
+
+    LimboListenQueueEntry(DocumentKey key, TargetData targetData) {
+      this.key = key;
+      this.targetData = targetData;
+    }
+
+    @NonNull
+    @Override
+    public String toString() {
+      return "key=" + key + " targetData=" + targetData;
+    }
+  }
+
   private static final String TAG = SyncEngine.class.getSimpleName();
 
+  private static final int DEFAULT_MAX_CONCURRENT_LIMBO_RESOLUTIONS = 100;
+
   /** Interface implemented by EventManager to handle notifications from SyncEngine. */
   interface SyncEngineCallback {
     /** Handles new view snapshots. */
@@ -109,6 +132,9 @@ interface SyncEngineCallback {
   /** The remote store for sending writes, watches, etc. to the backend. */
   private final RemoteStore remoteStore;
 
+  /** The max number of concurrent listens for limbo resolution; see {@link #limboListenQueue}. */
+  private final int maxConcurrentLimboResolutions;
+
   /** QueryViews for all active queries, indexed by query. */
   private final Map<Query, QueryView> queryViewsByQuery;
 
@@ -127,6 +153,13 @@ interface SyncEngineCallback {
    */
   private final Map<Integer, LimboResolution> limboResolutionsByTarget;
 
+  /**
+   * The list of enqueued limbo resolutions. After {@link #maxConcurrentLimboResolutions} listens
+   * have started for the purpose of limbo resolutions then additional limbo resolution listens are
+   * enqueued in this queue. When a listen completes then a listen is dequeued and started.
+   */
+  private final Queue<LimboListenQueueEntry> limboListenQueue;
+
   /** Used to track any documents that are currently in limbo. */
   private final ReferenceSet limboDocumentRefs;
 
@@ -144,14 +177,24 @@ interface SyncEngineCallback {
   private SyncEngineCallback syncEngineListener;
 
   public SyncEngine(LocalStore localStore, RemoteStore remoteStore, User initialUser) {
+    this(localStore, remoteStore, initialUser, DEFAULT_MAX_CONCURRENT_LIMBO_RESOLUTIONS);
+  }
+
+  public SyncEngine(
+      LocalStore localStore,
+      RemoteStore remoteStore,
+      User initialUser,
+      int maxConcurrentLimboResolutions) {
     this.localStore = localStore;
     this.remoteStore = remoteStore;
+    this.maxConcurrentLimboResolutions = maxConcurrentLimboResolutions;
 
     queryViewsByQuery = new HashMap<>();
     queriesByTarget = new HashMap<>();
 
     limboTargetsByKey = new HashMap<>();
     limboResolutionsByTarget = new HashMap<>();
+    limboListenQueue = new ArrayDeque<>();
     limboDocumentRefs = new ReferenceSet();
 
     mutationUserCallbacks = new HashMap<>();
@@ -380,6 +423,10 @@ public void handleRejectedListen(int targetId, Status error) {
       limboTargetsByKey.remove(limboKey);
       limboResolutionsByTarget.remove(targetId);
 
+      if (!limboListenQueue.isEmpty()) {
+        remoteStore.listen(limboListenQueue.remove().targetData);
+      }
+
       // TODO: Retry on transient errors?
 
       // It's a limbo doc. Create a synthetic event saying it was deleted. This is kind of a hack.
@@ -540,6 +587,9 @@ private void removeLimboTarget(DocumentKey key) {
       remoteStore.stopListening(targetId);
       limboTargetsByKey.remove(key);
       limboResolutionsByTarget.remove(targetId);
+      if (!limboListenQueue.isEmpty()) {
+        remoteStore.listen(limboListenQueue.remove().targetData);
+      }
     }
   }
 
@@ -616,7 +666,11 @@ private void trackLimboChange(LimboDocumentChange change) {
               ListenSequence.INVALID,
               QueryPurpose.LIMBO_RESOLUTION);
       limboResolutionsByTarget.put(limboTargetId, new LimboResolution(key));
-      remoteStore.listen(targetData);
+      if (limboResolutionsByTarget.size() <= maxConcurrentLimboResolutions) {
+        remoteStore.listen(targetData);
+      } else {
+        limboListenQueue.add(new LimboListenQueueEntry(key, targetData));
+      }
       limboTargetsByKey.put(key, limboTargetId);
     }
   }
@@ -627,6 +681,15 @@ public Map<DocumentKey, Integer> getCurrentLimboDocuments() {
     return new HashMap<>(limboTargetsByKey);
   }
 
+  @VisibleForTesting
+  public List<DocumentKey> getEnqueuedLimboDocuments() {
+    ArrayList<DocumentKey> list = new ArrayList<>(limboListenQueue.size());
+    for (LimboListenQueueEntry entry : limboListenQueue) {
+      list.add(entry.key);
+    }
+    return list;
+  }
+
   public void handleCredentialChange(User user) {
     boolean userChanged = !currentUser.equals(user);
     currentUser = user;

firebase-firestore/src/test/java/com/google/firebase/firestore/spec/SpecTestCase.java

@@ -153,6 +153,7 @@ public abstract class SpecTestCase implements RemoteStoreCallback {
           : Sets.newHashSet("no-android", BENCHMARK_TAG, "multi-client");
 
   private boolean garbageCollectionEnabled;
+  private Integer maxNumConcurrentLimboResolutions;
   private boolean networkEnabled = true;
 
   //
@@ -250,6 +251,8 @@ protected void specSetUp(JSONObject config) {
     outstandingWrites = new HashMap<>();
 
     this.garbageCollectionEnabled = config.optBoolean("useGarbageCollection", false);
+    this.maxNumConcurrentLimboResolutions =
+        (Integer) config.opt("maxNumConcurrentLimboResolutions");
 
     currentUser = User.UNAUTHENTICATED;
 
@@ -293,7 +296,12 @@ private void initClient() {
     ConnectivityMonitor connectivityMonitor =
         new AndroidConnectivityMonitor(ApplicationProvider.getApplicationContext());
     remoteStore = new RemoteStore(this, localStore, datastore, queue, connectivityMonitor);
-    syncEngine = new SyncEngine(localStore, remoteStore, currentUser);
+    if (maxNumConcurrentLimboResolutions == null) {
+      syncEngine = new SyncEngine(localStore, remoteStore, currentUser);
+    } else {
+      syncEngine =
+          new SyncEngine(localStore, remoteStore, currentUser, maxNumConcurrentLimboResolutions);
+    }
     eventManager = new EventManager(syncEngine);
     localStore.start();
     remoteStore.start();
@@ -988,12 +996,20 @@ private void validateLimboDocs() {
     @SuppressWarnings("VisibleForTests")
     Map<DocumentKey, Integer> actualLimboDocs =
         new HashMap<>(syncEngine.getCurrentLimboDocuments());
+    @SuppressWarnings("VisibleForTests")
+    List<DocumentKey> enqueuedLimboDocuments = syncEngine.getEnqueuedLimboDocuments();
 
     // Validate that each limbo doc has an expected active target
     for (Map.Entry<DocumentKey, Integer> limboDoc : actualLimboDocs.entrySet()) {
-      assertTrue(
-          "Found limbo doc " + limboDoc.getKey() + " without an expected active target",
-          expectedActiveTargets.containsKey(limboDoc.getValue()));
+      if (enqueuedLimboDocuments.contains(limboDoc.getKey())) {
+        assertFalse(
+            "Found limbo doc " + limboDoc.getKey() + " with an unexpected active target",
+            expectedActiveTargets.containsKey(limboDoc.getValue()));
+      } else {
+        assertTrue(
+            "Found limbo doc " + limboDoc.getKey() + " without an expected active target",
+            expectedActiveTargets.containsKey(limboDoc.getValue()));
+      }
     }
 
     for (DocumentKey expectedLimboDoc : expectedLimboDocs) {