awaitPendingWrites initial revision

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/FirestoreTest.java

@@ -1096,4 +1096,25 @@ public void testShutdownCalledMultipleTimes() {
 
     expectError(() -> waitFor(reference.get()), expectedMessage);
   }
+
+  @Test
+  public void testAwaitPendingWritesResolves() {
+    DocumentReference documentReference = testCollection("abc").document("123");
+    FirebaseFirestore firestore = documentReference.getFirestore();
+    Map<String, Object> data = map("foo", "bar");
+
+    waitFor(firestore.disableNetwork());
+    Task<Void> awaitsPendingWrites1 = firestore.awaitPendingWrites();
+    Task<Void> pendingWrite = documentReference.set(data);
+    Task<Void> awaitsPendingWrites2 = firestore.awaitPendingWrites();
+
+    assertTrue(!awaitsPendingWrites1.isComplete());
+    assertTrue(!pendingWrite.isComplete());
+    assertTrue(!awaitsPendingWrites2.isComplete());
+
+    waitFor(firestore.enableNetwork());
+    waitFor(pendingWrite);
+    waitFor(awaitsPendingWrites1);
+    waitFor(awaitsPendingWrites2);
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/FirebaseFirestore.java

@@ -362,6 +362,18 @@ Task<Void> shutdown() {
     return shutdownInternal();
   }
 
+  /**
+   * Wait for server acknowledgement for all pending writes existing at the time of calling this
+   * method.
+   *
+   * <p>Both acceptance and rejection count as server acknowledgement.
+   *
+   * @return A {@link Task} which resolves when all pending writes are acknowledged by the server.
+   */
+  Task<Void> awaitPendingWrites() {
+    return client.awaitPendingWrites();
+  }
+
   @VisibleForTesting
   AsyncQueue getAsyncQueue() {
     return asyncQueue;

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

@@ -224,6 +224,24 @@ public <TResult> Task<TResult> transaction(
         () -> syncEngine.transaction(asyncQueue, updateFunction, retries));
   }
 
+  /**
+   * Returns a task resolves when all the pending writes at the time when this method is called
+   * received server acknowledgement. An acknowledgement can be either acceptance or rejections.
+   */
+  public Task<Void> awaitPendingWrites() {
+    this.verifyNotShutdown();
+    if (!remoteStore.canUseNetwork()) {
+      Logger.warn(
+          LOG_TAG,
+          "Network is disabled, the Task created to wait for all writes getting"
+              + " acknowledged by server will not complete until network is enabled.");
+    }
+
+    final TaskCompletionSource<Void> source = new TaskCompletionSource<>();
+    asyncQueue.enqueueAndForget(() -> syncEngine.registerPendingWritesTask(source));
+    return source.getTask();
+  }
+
   private void initialize(Context context, User user, boolean usePersistence, long cacheSizeBytes) {
     // Note: The initialization work must all be synchronous (we can't dispatch more work) since
     // external write/listen operations could get queued to run before that subsequent work

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

@@ -23,6 +23,7 @@
 import com.google.android.gms.tasks.TaskCompletionSource;
 import com.google.android.gms.tasks.Tasks;
 import com.google.common.base.Function;
+import com.google.common.collect.Lists;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.FirebaseFirestoreException;
@@ -133,6 +134,9 @@ interface SyncEngineCallback {
   /** Stores user completion blocks, indexed by user and batch ID. */
   private final Map<User, Map<Integer, TaskCompletionSource<Void>>> mutationUserCallbacks;
 
+  /** Stores user callbacks waiting for all pending writes to be acknowledged. */
+  private final Map<Integer, List<TaskCompletionSource<Void>>> pendingWritesCallbacks;
+
   /** Used for creating the target IDs for the listens used to resolve limbo documents. */
   private final TargetIdGenerator targetIdGenerator;
 
@@ -154,6 +158,8 @@ public SyncEngine(LocalStore localStore, RemoteStore remoteStore, User initialUs
     mutationUserCallbacks = new HashMap<>();
     targetIdGenerator = TargetIdGenerator.forSyncEngine();
     currentUser = initialUser;
+
+    pendingWritesCallbacks = new HashMap<>();
   }
 
   public void setCallback(SyncEngineCallback callback) {
@@ -407,6 +413,8 @@ public void handleSuccessfulWrite(MutationBatchResult mutationBatchResult) {
     // they consistently happen before listen events.
     notifyUser(mutationBatchResult.getBatch().getBatchId(), /*status=*/ null);
 
+    resolveTasksAwaitingForPendingWritesIfAny(mutationBatchResult.getBatch().getBatchId());
+
     ImmutableSortedMap<DocumentKey, MaybeDocument> changes =
         localStore.acknowledgeBatch(mutationBatchResult);
 
@@ -427,9 +435,41 @@ public void handleRejectedWrite(int batchId, Status status) {
     // they consistently happen before listen events.
     notifyUser(batchId, status);
 
+    resolveTasksAwaitingForPendingWritesIfAny(batchId);
+
     emitNewSnapsAndNotifyLocalStore(changes, /*remoteEvent=*/ null);
   }
 
+  /**
+   * Takes a snapshot of current local mutation queue, and register a user task which will resolve
+   * when all those mutations are either accepted or rejected by the server.
+   */
+  public void registerPendingWritesTask(TaskCompletionSource<Void> userTask) {
+    int largestPendingBatchId = localStore.getHighestUnacknowledgedBatchId();
+
+    if (largestPendingBatchId == 0) {
+      // Complete the task right away if there is no pending writes at the moment.
+      userTask.setResult(null);
+    }
+
+    if (pendingWritesCallbacks.containsKey(largestPendingBatchId)) {
+      pendingWritesCallbacks.get(largestPendingBatchId).add(userTask);
+    } else {
+      pendingWritesCallbacks.put(largestPendingBatchId, Lists.newArrayList(userTask));
+    }
+  }
+
+  /** Resolves tasks waiting for this batch id to get acknowledged by server, if there is any. */
+  private void resolveTasksAwaitingForPendingWritesIfAny(int batchId) {
+    if (pendingWritesCallbacks.containsKey(batchId)) {
+      for (TaskCompletionSource<Void> task : pendingWritesCallbacks.get(batchId)) {
+        task.setResult(null);
+      }
+
+      pendingWritesCallbacks.remove(batchId);
+    }
+  }
+
   /** Resolves the task corresponding to this write result. */
   private void notifyUser(int batchId, @Nullable Status status) {
     Map<Integer, TaskCompletionSource<Void>> userTasks = mutationUserCallbacks.get(currentUser);

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

@@ -282,6 +282,14 @@ public ImmutableSortedMap<DocumentKey, MaybeDocument> rejectBatch(int batchId) {
         });
   }
 
+  /**
+   * Returns the largest (latest) batch id in mutation queue that is pending server response.
+   * Returns 0 if the queue is empty.
+   */
+  public int getHighestUnacknowledgedBatchId() {
+    return mutationQueue.getLargestUnacknowledgedBatchId();
+  }
+
   /** Returns the last recorded stream token for the current user. */
   public ByteString getLastStreamToken() {
     return mutationQueue.getLastStreamToken();

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

@@ -187,6 +187,11 @@ public MutationBatch getNextMutationBatchAfterBatchId(int batchId) {
     return queue.size() > index ? queue.get(index) : null;
   }
 
+  @Override
+  public int getLargestUnacknowledgedBatchId() {
+    return queue.size() == 0 ? 0 : nextBatchId - 1;
+  }
+
   @Override
   public List<MutationBatch> getAllMutationBatches() {
     return Collections.unmodifiableList(queue);

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

@@ -73,6 +73,12 @@ MutationBatch addMutationBatch(
   @Nullable
   MutationBatch getNextMutationBatchAfterBatchId(int batchId);
 
+  /**
+   * @return The largest (latest) batch id in mutation queue for the current user that is pending
+   *     server response, 0 if the queue is empty.
+   */
+  int getLargestUnacknowledgedBatchId();
+
   /** Returns all mutation batches in the mutation queue. */
   // TODO: PERF: Current consumer only needs mutated keys; if we can provide that
   // cheaply, we should replace this.

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

@@ -249,6 +249,16 @@ public MutationBatch getNextMutationBatchAfterBatchId(int batchId) {
         .firstValue(row -> decodeInlineMutationBatch(row.getInt(0), row.getBlob(1)));
   }
 
+  @Override
+  public int getLargestUnacknowledgedBatchId() {
+    if (isEmpty()) {
+      return 0;
+    }
+    return db.query("SELECT MAX(batch_id) FROM mutations " + "WHERE uid = ?")
+        .binding(uid)
+        .firstValue(row -> row.getInt(0));
+  }
+
   @Override
   public List<MutationBatch> getAllMutationBatches() {
     List<MutationBatch> result = new ArrayList<>();

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

@@ -478,7 +478,7 @@ private void handleWatchStreamClose(Status status) {
     }
   }
 
-  private boolean canUseNetwork() {
+  public boolean canUseNetwork() {
     // PORTING NOTE: This method exists mostly because web also has to take into account primary
     // vs. secondary state.
     return networkEnabled;

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java

@@ -1143,4 +1143,21 @@ public void testHandlesPatchMutationWithTransformThenRemoteEvent() {
     assertChanged(doc("foo/bar", 1, map("sum", 1), Document.DocumentState.LOCAL_MUTATIONS));
     assertContains(doc("foo/bar", 1, map("sum", 1), Document.DocumentState.LOCAL_MUTATIONS));
   }
+
+  @Test
+  public void testGetHighestUnacknowledgedBatchIdReturnsExpectedResult() {
+    assertEquals(0, localStore.getHighestUnacknowledgedBatchId());
+
+    writeMutation(setMutation("foo/bar", map("abc", 123)));
+    assertEquals(1, localStore.getHighestUnacknowledgedBatchId());
+
+    writeMutation(patchMutation("foo/bar", map("abc", 321)));
+    assertEquals(2, localStore.getHighestUnacknowledgedBatchId());
+
+    acknowledgeMutation(1);
+    assertEquals(2, localStore.getHighestUnacknowledgedBatchId());
+
+    rejectMutation();
+    assertEquals(0, localStore.getHighestUnacknowledgedBatchId());
+  }
 }