Port waitForPendingWrites (#3619)

* port waitForPendingWrites

* addressing comments #1

* remove commented line

* another change to disable network

Author: wu-hui

Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm

@@ -1402,4 +1402,50 @@ - (void)testCanRemoveListenerAfterShutdown {
   [listenerRegistration remove];
 }
 
+- (void)testWaitForPendingWritesCompletes {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+
+  [self disableNetwork];
+
+  [doc setData:@{@"foo" : @"bar"}];
+  [firestore waitForPendingWritesWithCompletion:
+                 [self completionForExpectationWithName:@"Wait for pending writes"]];
+
+  [firestore enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable network"]];
+  [self awaitExpectations];
+}
+
+- (void)testWaitForPendingWritesFailsWhenUserChanges {
+  FIRFirestore *firestore = self.db;
+
+  [self disableNetwork];
+
+  // Writes to local to prevent immediate call to the completion of waitForPendingWrites.
+  NSDictionary<NSString *, id> *data =
+      @{@"owner" : @{@"name" : @"Andy", @"email" : @"[email protected]"}};
+  [[self documentRef] setData:data];
+
+  XCTestExpectation *expectation = [self expectationWithDescription:@"waitForPendingWrites"];
+  [firestore waitForPendingWritesWithCompletion:^(NSError *_Nullable error) {
+    XCTAssertNotNil(error);
+    XCTAssertEqualObjects(error.domain, FIRFirestoreErrorDomain);
+    XCTAssertEqual(error.code, FIRFirestoreErrorCodeCancelled);
+    [expectation fulfill];
+  }];
+
+  [self triggerUserChangeWithUid:@"user-to-fail-pending-writes"];
+  [self awaitExpectations];
+}
+
+- (void)testWaitForPendingWritesCompletesWhenOfflineIfNoPending {
+  FIRFirestore *firestore = self.db;
+
+  [self disableNetwork];
+
+  [firestore waitForPendingWritesWithCompletion:
+                 [self completionForExpectationWithName:@"Wait for pending writes"]];
+  [self awaitExpectations];
+}
+
 @end

Firestore/Example/Tests/Local/FSTLevelDBMutationQueueTests.mm

@@ -87,14 +87,14 @@ - (void)setUp {
 
 - (void)testLoadNextBatchID_zeroWhenTotallyEmpty {
   // Initial seek is invalid
-  XCTAssertEqual(LoadNextBatchIdFromDb(_db.ptr), 0);
+  XCTAssertEqual(LoadNextBatchIdFromDb(_db.ptr), 1);
 }
 
 - (void)testLoadNextBatchID_zeroWhenNoMutations {
   // Initial seek finds no mutations
   [self setDummyValueForKey:MutationLikeKey("mutationr", "foo", 20)];
   [self setDummyValueForKey:MutationLikeKey("mutationsa", "foo", 10)];
-  XCTAssertEqual(LoadNextBatchIdFromDb(_db.ptr), 0);
+  XCTAssertEqual(LoadNextBatchIdFromDb(_db.ptr), 1);
 }
 
 - (void)testLoadNextBatchID_findsSingleRow {

Firestore/Example/Tests/Local/FSTLocalStoreTests.mm

@@ -1105,6 +1105,24 @@ - (void)testHandlesPatchMutationWithTransformThenRemoteEvent {
   FSTAssertChanged(Doc("foo/bar", 1, Map("sum", 1), DocumentState::kLocalMutations));
 }
 
+- (void)testGetHighestUnacknowledgeBatchId {
+  if ([self isTestBaseClass]) return;
+
+  XCTAssertEqual(-1, [self.localStore getHighestUnacknowledgedBatchId]);
+
+  [self writeMutation:FSTTestSetMutation(@"foo/bar", @{@"abc" : @123})];
+  XCTAssertEqual(1, [self.localStore getHighestUnacknowledgedBatchId]);
+
+  [self writeMutation:FSTTestPatchMutation("foo/bar", @{@"abc" : @321}, {})];
+  XCTAssertEqual(2, [self.localStore getHighestUnacknowledgedBatchId]);
+
+  [self acknowledgeMutationWithVersion:1];
+  XCTAssertEqual(2, [self.localStore getHighestUnacknowledgedBatchId]);
+
+  [self rejectMutation];
+  XCTAssertEqual(-1, [self.localStore getHighestUnacknowledgedBatchId]);
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Example/Tests/Util/FSTIntegrationTestCase.h

@@ -55,7 +55,12 @@ extern "C" {
 /** Returns a new Firestore connected to the project with the given projectID. */
 - (FIRFirestore *)firestoreWithProjectID:(NSString *)projectID;
 
-/** Returns a new Firestore connected to the project with the given app. */
+/** Triggers a user change with given user id. */
+- (void)triggerUserChangeWithUid:(NSString *)uid;
+
+/**
+ * Returns a new Firestore connected to the project with the given app.
+ */
 - (FIRFirestore *)firestoreWithApp:(FIRApp *)app;
 
 /** Synchronously shuts down the given firestore. */

Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm

@@ -33,7 +33,9 @@
 #include <string>
 #include <utility>
 
+#include "Firestore/core/src/firebase/firestore/auth/credentials_provider.h"
 #include "Firestore/core/src/firebase/firestore/auth/empty_credentials_provider.h"
+#include "Firestore/core/src/firebase/firestore/auth/user.h"
 #include "Firestore/core/src/firebase/firestore/model/database_id.h"
 #include "Firestore/core/src/firebase/firestore/remote/grpc_connection.h"
 #include "Firestore/core/src/firebase/firestore/util/autoid.h"
@@ -55,8 +57,10 @@
 #include "Firestore/core/src/firebase/firestore/util/executor_libdispatch.h"
 
 namespace util = firebase::firestore::util;
+using firebase::firestore::auth::CredentialChangeListener;
 using firebase::firestore::auth::CredentialsProvider;
 using firebase::firestore::auth::EmptyCredentialsProvider;
+using firebase::firestore::auth::User;
 using firebase::firestore::model::DatabaseId;
 using firebase::firestore::testutil::AppForUnitTesting;
 using firebase::firestore::remote::GrpcConnection;
@@ -79,13 +83,37 @@
 
 static bool runningAgainstEmulator = false;
 
+// Behaves the same as `EmptyCredentialsProvider` except it can also trigger a user
+// change.
+class FakeCredentialsProvider : public EmptyCredentialsProvider {
+ public:
+  void SetCredentialChangeListener(CredentialChangeListener changeListener) override {
+    if (changeListener) {
+      listener_ = std::move(changeListener);
+      listener_(User::Unauthenticated());
+    }
+  }
+
+  void ChangeUser(NSString *new_id) {
+    if (listener_) {
+      listener_(firebase::firestore::auth::User::FromUid(new_id));
+    }
+  }
+
+ private:
+  CredentialChangeListener listener_;
+};
+
 @implementation FSTIntegrationTestCase {
   NSMutableArray<FIRFirestore *> *_firestores;
+  std::shared_ptr<FakeCredentialsProvider> _fakeCredentialsProvider;
 }
 
 - (void)setUp {
   [super setUp];
 
+  _fakeCredentialsProvider = std::make_shared<FakeCredentialsProvider>();
+
   [self clearPersistenceOnce];
   [self primeBackend];
 
@@ -239,13 +267,11 @@ - (FIRFirestore *)firestoreWithApp:(FIRApp *)app {
 
   FIRSetLoggerLevel(FIRLoggerLevelDebug);
 
-  std::unique_ptr<CredentialsProvider> credentials_provider =
-      absl::make_unique<firebase::firestore::auth::EmptyCredentialsProvider>();
   std::string projectID = util::MakeString(app.options.projectID);
   FIRFirestore *firestore =
       [[FIRFirestore alloc] initWithDatabaseID:DatabaseId(projectID)
                                 persistenceKey:util::MakeString(persistenceKey)
-                           credentialsProvider:std::move(credentials_provider)
+                           credentialsProvider:_fakeCredentialsProvider
                                    workerQueue:std::move(workerQueue)
                                    firebaseApp:app
                               instanceRegistry:nil];
@@ -255,6 +281,10 @@ - (FIRFirestore *)firestoreWithApp:(FIRApp *)app {
   return firestore;
 }
 
+- (void)triggerUserChangeWithUid:(NSString *)uid {
+  _fakeCredentialsProvider->ChangeUser(uid);
+}
+
 - (void)primeBackend {
   static dispatch_once_t onceToken;
   dispatch_once(&onceToken, ^{

Firestore/Source/API/FIRFirestore+Internal.h

@@ -90,6 +90,22 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)shutdownInternalWithCompletion:(nullable void (^)(NSError *_Nullable error))completion;
 
+/**
+ * Waits until all currently pending writes for the active user have been acknowledged by the
+ * backend.
+ *
+ * The completion block is called immediately without error if there are no outstanding writes.
+ * Otherwise, the completion block is called when all previously issued writes (including those
+ * written in a previous app session) have been acknowledged by the backend. The completion
+ * block does not wait for writes that were added after the method is called. If you
+ * wish to wait for additional writes, you have to call `waitForPendingWritesWithCompletion`
+ * again.
+ *
+ * Any outstanding `waitForPendingWritesWithCompletion` completion blocks are called with an
+ * error during user change.
+ */
+- (void)waitForPendingWritesWithCompletion:(void (^)(NSError *_Nullable error))completion;
+
 - (const std::shared_ptr<util::AsyncQueue> &)workerQueue;
 
 @property(nonatomic, assign, readonly) std::shared_ptr<api::Firestore> wrapped;

Firestore/Source/API/FIRFirestore.mm

@@ -194,6 +194,10 @@ - (FIRWriteBatch *)batch {
                                          writeBatch:_firestore->GetBatch()];
 }
 
+- (void)waitForPendingWritesWithCompletion:(void (^)(NSError *_Nullable error))completion {
+  _firestore->WaitForPendingWrites(util::MakeCallback(completion));
+}
+
 - (void)runTransactionWithBlock:(id _Nullable (^)(FIRTransaction *, NSError **))updateBlock
                   dispatchQueue:(dispatch_queue_t)queue
                      completion:

Firestore/Source/Core/FSTFirestoreClient.h

@@ -111,6 +111,13 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)writeMutations:(std::vector<model::Mutation> &&)mutations
               callback:(util::StatusCallback)callback;
 
+/**
+ * Passes a callback that is triggered when all the pending writes at the
+ * time when this method is called received server acknowledgement.
+ * An acknowledgement can be either acceptance or rejections.
+ */
+- (void)waitForPendingWritesWithCallback:(util::StatusCallback)callback;
+
 /** Tries to execute the transaction in updateCallback up to retries times. */
 - (void)transactionWithRetries:(int)retries
                 updateCallback:(core::TransactionUpdateCallback)updateCallback

Firestore/Source/Core/FSTFirestoreClient.mm

@@ -442,6 +442,20 @@ - (void)writeMutations:(std::vector<Mutation> &&)mutations callback:(util::Statu
   });
 };
 
+- (void)waitForPendingWritesWithCallback:(util::StatusCallback)callback {
+  [self verifyNotShutdown];
+  // Dispatch the result back onto the user dispatch queue.
+  auto async_callback = [self, callback](util::Status status) {
+    if (callback) {
+      self->_userExecutor->Execute([=] { callback(std::move(status)); });
+    }
+  };
+
+  _workerQueue->Enqueue([self, async_callback]() {
+    [self.syncEngine registerPendingWritesCallback:std::move(async_callback)];
+  });
+}
+
 - (void)transactionWithRetries:(int)retries
                 updateCallback:(core::TransactionUpdateCallback)update_callback
                 resultCallback:(core::TransactionResultCallback)resultCallback {

Firestore/Source/Core/FSTSyncEngine.h

@@ -86,6 +86,12 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)writeMutations:(std::vector<model::Mutation> &&)mutations
             completion:(FSTVoidErrorBlock)completion;
 
+/**
+ * Registers a user callback that is called when all pending mutations at the moment of calling
+ * are acknowledged .
+ */
+- (void)registerPendingWritesCallback:(util::StatusCallback)callback;
+
 /**
  * Runs the given transaction block up to retries times and then calls completion.
  *

Firestore/Source/Core/FSTSyncEngine.mm

@@ -40,6 +40,7 @@
 #include "Firestore/core/src/firebase/firestore/model/document_key.h"
 #include "Firestore/core/src/firebase/firestore/model/document_map.h"
 #include "Firestore/core/src/firebase/firestore/model/document_set.h"
+#include "Firestore/core/src/firebase/firestore/model/mutation_batch.h"
 #include "Firestore/core/src/firebase/firestore/model/no_document.h"
 #include "Firestore/core/src/firebase/firestore/model/snapshot_version.h"
 #include "Firestore/core/src/firebase/firestore/remote/remote_event.h"
@@ -64,6 +65,7 @@
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::DocumentKeySet;
 using firebase::firestore::model::DocumentMap;
+using firebase::firestore::model::kBatchIdUnknown;
 using firebase::firestore::model::ListenSequenceNumber;
 using firebase::firestore::model::MaybeDocumentMap;
 using firebase::firestore::model::Mutation;
@@ -79,6 +81,7 @@
 using firebase::firestore::util::AsyncQueue;
 using firebase::firestore::util::MakeNSError;
 using firebase::firestore::util::Status;
+using firebase::firestore::util::StatusCallback;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -197,6 +200,9 @@ @implementation FSTSyncEngine {
   std::unordered_map<User, NSMutableDictionary<NSNumber *, FSTVoidErrorBlock> *, HashUser>
       _mutationCompletionBlocks;
 
+  /** Stores user callbacks waiting for pending writes to be acknowledged. */
+  std::unordered_map<model::BatchId, std::vector<StatusCallback>> _pendingWritesCallbacks;
+
   /** FSTQueryViews for all active queries, indexed by query. */
   std::unordered_map<Query, FSTQueryView *> _queryViewsByQuery;
 
@@ -296,6 +302,51 @@ - (void)writeMutations:(std::vector<Mutation> &&)mutations
   _remoteStore->FillWritePipeline();
 }
 
+- (void)registerPendingWritesCallback:(StatusCallback)callback {
+  if (!_remoteStore->CanUseNetwork()) {
+    LOG_DEBUG("The network is disabled. The task returned by 'awaitPendingWrites()' will not "
+              "complete until the network is enabled.");
+  }
+
+  int largestPendingBatchId = [self.localStore getHighestUnacknowledgedBatchId];
+
+  if (largestPendingBatchId == kBatchIdUnknown) {
+    // Trigger the callback right away if there is no pending writes at the moment.
+    callback(Status::OK());
+    return;
+  }
+
+  auto it = _pendingWritesCallbacks.find(largestPendingBatchId);
+  if (it != _pendingWritesCallbacks.end()) {
+    it->second.push_back(std::move(callback));
+  } else {
+    _pendingWritesCallbacks.emplace(largestPendingBatchId,
+                                    std::vector<StatusCallback>({std::move(callback)}));
+  }
+}
+
+/** Triggers callbacks waiting for this batch id to get acknowledged by server, if there are any. */
+- (void)triggerPendingWriteCallbacksWithBatchId:(int)batchId {
+  auto it = _pendingWritesCallbacks.find(batchId);
+  if (it != _pendingWritesCallbacks.end()) {
+    for (const auto &callback : it->second) {
+      callback(Status::OK());
+    }
+
+    _pendingWritesCallbacks.erase(it);
+  }
+}
+
+- (void)failOutstandingPendingWritesAwaitingCallbacks:(absl::string_view)errorMessage {
+  for (const auto &entry : _pendingWritesCallbacks) {
+    for (const auto &callback : entry.second) {
+      callback(Status(Error::Cancelled, errorMessage));
+    }
+  }
+
+  _pendingWritesCallbacks.clear();
+}
+
 - (void)addMutationCompletionBlock:(FSTVoidErrorBlock)completion batchID:(BatchId)batchID {
   NSMutableDictionary<NSNumber *, FSTVoidErrorBlock> *completionBlocks =
       _mutationCompletionBlocks[_currentUser];
@@ -460,6 +511,8 @@ - (void)applySuccessfulWriteWithResult:(FSTMutationBatchResult *)batchResult {
   // consistently happen before listen events.
   [self processUserCallbacksForBatchID:batchResult.batch.batchID error:nil];
 
+  [self triggerPendingWriteCallbacksWithBatchId:batchResult.batch.batchID];
+
   MaybeDocumentMap changes = [self.localStore acknowledgeBatchWithResult:batchResult];
   [self emitNewSnapshotsAndNotifyLocalStoreWithChanges:changes remoteEvent:absl::nullopt];
 }
@@ -478,6 +531,8 @@ - (void)rejectFailedWriteWithBatchID:(BatchId)batchID error:(NSError *)error {
   // consistently happen before listen events.
   [self processUserCallbacksForBatchID:batchID error:error];
 
+  [self triggerPendingWriteCallbacksWithBatchId:batchID];
+
   [self emitNewSnapshotsAndNotifyLocalStoreWithChanges:changes remoteEvent:absl::nullopt];
 }
 
@@ -629,6 +684,9 @@ - (void)credentialDidChangeWithUser:(const firebase::firestore::auth::User &)use
   _currentUser = user;
 
   if (userChanged) {
+    // Fails callbacks waiting for pending writes requested by previous user.
+    [self failOutstandingPendingWritesAwaitingCallbacks:
+              "'waitForPendingWrites' callback is cancelled due to a user change."];
     // Notify local store and emit any resulting events from swapping out the mutation queue.
     MaybeDocumentMap changes = [self.localStore userDidChange:user];
     [self emitNewSnapshotsAndNotifyLocalStoreWithChanges:changes remoteEvent:absl::nullopt];