awaitPendingWrites initial revision

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

@@ -17,6 +17,7 @@
 import static com.google.firebase.firestore.AccessHelper.getAsyncQueue;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.newTestSettings;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.provider;
+import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testChangeUserTo;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollection;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testCollectionWithDocs;
 import static com.google.firebase.firestore.testutil.IntegrationTestUtil.testDocument;
@@ -44,6 +45,7 @@
 import com.google.firebase.Timestamp;
 import com.google.firebase.firestore.FirebaseFirestoreException.Code;
 import com.google.firebase.firestore.Query.Direction;
+import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.testutil.EventAccumulator;
 import com.google.firebase.firestore.testutil.IntegrationTestUtil;
 import com.google.firebase.firestore.util.AsyncQueue.TimerId;
@@ -1096,4 +1098,49 @@ public void testShutdownCalledMultipleTimes() {
 
     expectError(() -> waitFor(reference.get()), expectedMessage);
   }
+
+  @Test
+  public void testWaitForPendingWritesResolves() {
+    DocumentReference documentReference = testCollection("abc").document("123");
+    FirebaseFirestore firestore = documentReference.getFirestore();
+    Map<String, Object> data = map("foo", "bar");
+
+    waitFor(firestore.disableNetwork());
+    Task<Void> awaitsPendingWrites1 = firestore.waitForPendingWrites();
+    Task<Void> pendingWrite = documentReference.set(data);
+    Task<Void> awaitsPendingWrites2 = firestore.waitForPendingWrites();
+
+    // `awaitsPendingWrites1` is complete immediately because there is no pending writes at
+    // the time it is created.
+    waitFor(awaitsPendingWrites1);
+    assertTrue(awaitsPendingWrites1.isComplete() && awaitsPendingWrites1.isSuccessful());
+    assertTrue(!pendingWrite.isComplete());
+    assertTrue(!awaitsPendingWrites2.isComplete());
+
+    waitFor(firestore.enableNetwork());
+    waitFor(pendingWrite);
+    waitFor(awaitsPendingWrites2);
+    assertTrue(awaitsPendingWrites2.isComplete() && awaitsPendingWrites2.isSuccessful());
+  }
+
+  @Test
+  public void testWaitForPendingWritesFailsWhenUserChanges() {
+    DocumentReference documentReference = testCollection("abc").document("123");
+    FirebaseFirestore firestore = documentReference.getFirestore();
+    Map<String, Object> data = map("foo", "bar");
+
+    // Prevent pending writes receiving acknowledgement.
+    waitFor(firestore.disableNetwork());
+    Task<Void> pendingWrite = documentReference.set(data);
+    Task<Void> awaitsPendingWrites = firestore.waitForPendingWrites();
+    assertTrue(!pendingWrite.isComplete());
+    assertTrue(!awaitsPendingWrites.isComplete());
+
+    testChangeUserTo(new User("new user"));
+
+    assertTrue(!pendingWrite.isComplete());
+    assertEquals(
+        "'waitForPendingWrites' task is cancelled due to User change.",
+        waitForException(awaitsPendingWrites).getMessage());
+  }
 }

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -36,11 +36,13 @@
 import com.google.firebase.firestore.MetadataChanges;
 import com.google.firebase.firestore.QuerySnapshot;
 import com.google.firebase.firestore.auth.EmptyCredentialsProvider;
+import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.DatabaseInfo;
 import com.google.firebase.firestore.local.Persistence;
 import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.testutil.provider.FirestoreProvider;
 import com.google.firebase.firestore.util.AsyncQueue;
+import com.google.firebase.firestore.util.Listener;
 import com.google.firebase.firestore.util.Logger;
 import com.google.firebase.firestore.util.Logger.Level;
 import java.util.ArrayList;
@@ -53,6 +55,34 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+class MockCredentialsProvider extends EmptyCredentialsProvider {
+
+  private static MockCredentialsProvider instance;
+
+  public static MockCredentialsProvider instance() {
+    if (MockCredentialsProvider.instance == null) {
+      MockCredentialsProvider.instance = new MockCredentialsProvider();
+    }
+    return MockCredentialsProvider.instance;
+  }
+
+  private MockCredentialsProvider() {
+    super();
+  }
+
+  @Override
+  public void setChangeListener(Listener<User> changeListener) {
+    super.setChangeListener(changeListener);
+    this.listener = changeListener;
+  }
+
+  public void changeUserTo(User user) {
+    listener.onValue(user);
+  }
+
+  private Listener<User> listener;
+}
+
 /** A set of helper methods for tests */
 public class IntegrationTestUtil {
 
@@ -239,7 +269,7 @@ public static FirebaseFirestore testFirestore(
             context,
             databaseId,
             persistenceKey,
-            new EmptyCredentialsProvider(),
+            MockCredentialsProvider.instance(),
             asyncQueue,
             /*firebaseApp=*/ null,
             /*instanceRegistry=*/ (dbId) -> {});
@@ -409,4 +439,8 @@ public static Map<String, Object> toDataMap(QuerySnapshot qrySnap) {
   public static boolean isRunningAgainstEmulator() {
     return CONNECT_TO_EMULATOR;
   }
+
+  public static void testChangeUserTo(User user) {
+    MockCredentialsProvider.instance().changeUserTo(user);
+  }
 }

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> waitForPendingWrites() {
+    return client.waitForPendingWrites();
+  }
+
   @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> waitForPendingWrites() {
+    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,56 @@ 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);
+      return;
+    }
+
+    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);
+    }
+  }
+
+  private void failOutstandingPendingWritesAwaitingTasks() {
+    for (Map.Entry<Integer, List<TaskCompletionSource<Void>>> entry :
+        pendingWritesCallbacks.entrySet()) {
+      for (TaskCompletionSource<Void> task : entry.getValue()) {
+        task.setException(
+            new FirebaseFirestoreException(
+                "'waitForPendingWrites' task is cancelled due to User change.",
+                FirebaseFirestoreException.Code.CANCELLED));
+      }
+    }
+
+    pendingWritesCallbacks.clear();
+  }
+
   /** Resolves the task corresponding to this write result. */
   private void notifyUser(int batchId, @Nullable Status status) {
     Map<Integer, TaskCompletionSource<Void>> userTasks = mutationUserCallbacks.get(currentUser);
@@ -565,6 +620,7 @@ public void handleCredentialChange(User user) {
       // Notify local store and emit any resulting events from swapping out the mutation queue.
       ImmutableSortedMap<DocumentKey, MaybeDocument> changes = localStore.handleUserChange(user);
       emitNewSnapsAndNotifyLocalStore(changes, /*remoteEvent=*/ null);
+      failOutstandingPendingWritesAwaitingTasks();
     }
 
     // Notify remote store so it can restart its streams.

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());
+  }
 }