Have ListenerRegistration::Remove block until callbacks are complete (#5961)

* Add the ability to await a single expectation

* Await individual expectations in test helpers

This prevents tests from accidentally awaiting unrelated expectations
when using the helpers.

* Fix tests that implicitly relied on test helpers to await all

* Add a test that verifies ListenerRegistration::Remove blocks

* Block ListenerRegistration::Remove when there are outstanding callbacks.

* Test that remove does not block when called from the callback.

Author: wilhuff

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

@@ -566,6 +566,7 @@ - (void)testSnapshotsInSyncListenerFiresAfterListenersInSync {
   }];
 
   [self writeDocumentRef:ref data:@{@"foo" : @3}];
+  [self awaitExpectation:done];
 }
 
 - (void)testSnapshotsInSyncRemoveIsIdempotent {
@@ -1484,6 +1485,104 @@ - (void)testCanRemoveListenerAfterTermination {
   [listenerRegistration remove];
 }
 
+- (void)testListenerCallbackBlocksRemove {
+  // This tests a guarantee required for C++ that doesn't strictly matter for Objective-C and has no
+  // equivalent on other platforms.
+  //
+  // The problem for C++ is that users can register a listener that refers to some state, then call
+  // `ListenerRegistration::Remove()` and expect to be able to immediately delete that state. The
+  // trouble is that there may be a callback in progress against that listener so the implementation
+  // now blocks the remove call until the callback is complete.
+  XCTestExpectation *running = [self expectationWithDescription:@"listener running"];
+  XCTestExpectation *allowCompletion =
+      [self expectationWithDescription:@"allow listener to complete"];
+  XCTestExpectation *removing = [self expectationWithDescription:@"attempting to remove listener"];
+  XCTestExpectation *removed = [self expectationWithDescription:@"listener removed"];
+
+  NSMutableString *steps = [NSMutableString string];
+
+  FIRDocumentReference *doc = [self documentRef];
+  [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
+
+  __block bool firstTime = true;
+
+  id<FIRListenerRegistration> listener =
+      [doc addSnapshotListener:^(FIRDocumentSnapshot *, NSError *) {
+        @synchronized(self) {
+          if (!firstTime) {
+            return;
+          }
+          firstTime = false;
+        }
+
+        [steps appendString:@"1"];
+        [running fulfill];
+
+        [self awaitExpectation:allowCompletion];
+        [steps appendString:@"3"];
+      }];
+
+  // Call remove asynchronously to avoid blocking the main test thread.
+  dispatch_queue_t async = dispatch_queue_create("firestore.async", DISPATCH_QUEUE_SERIAL);
+  dispatch_async(async, ^{
+    [self awaitExpectation:running];
+    [steps appendString:@"2"];
+
+    [removing fulfill];
+    [listener remove];
+
+    [steps appendString:@"4"];
+    [removed fulfill];
+  });
+
+  // Perform a write to `doc` which will trigger the listener callback. Don't wait for completion
+  // though because that completion handler is in line behind the listener callback that the test
+  // is blocking.
+  XCTestExpectation *setData = [self expectationWithDescription:@"setData"];
+  [doc setData:@{@"foo" : @"bar"} completion:[self completionForExpectation:setData]];
+
+  [self awaitExpectation:removing];
+  [allowCompletion fulfill];
+
+  [self awaitExpectation:removed];
+  XCTAssertEqualObjects(steps, @"1234");
+
+  [self awaitExpectation:setData];
+}
+
+- (void)testListenerCallbackCanCallRemoveWithoutBlocking {
+  // This tests a guarantee required for C++ that doesn't strictly matter for Objective-C and has no
+  // equivalent on other platforms. See `testListenerCallbackBlocksRemove` for background.
+  XCTestExpectation *removed = [self expectationWithDescription:@"listener removed"];
+
+  NSMutableString *steps = [NSMutableString string];
+
+  FIRDocumentReference *doc = [self documentRef];
+  [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
+
+  __block id<FIRListenerRegistration> listener = nil;
+
+  @synchronized(self) {
+    listener = [doc addSnapshotListener:^(FIRDocumentSnapshot *, NSError *) {
+      [steps appendString:@"1"];
+
+      @synchronized(self) {
+        // This test is successful if this method does not block.
+        [listener remove];
+      }
+
+      [steps appendString:@"2"];
+      [removed fulfill];
+    }];
+  }
+
+  // Perform a write to `doc` which will trigger the listener callback.
+  [self writeDocumentRef:doc data:@{@"foo" : @"bar2"}];
+
+  [self awaitExpectation:removed];
+  XCTAssertEqualObjects(steps, @"12");
+}
+
 - (void)testWaitForPendingWritesCompletes {
   FIRDocumentReference *doc = [self documentRef];
   FIRFirestore *firestore = doc.firestore;

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

@@ -158,6 +158,7 @@ - (void)testCanOutliveDocumentReference {
 
   FIRDocumentReference *docRef2 = [collectionRef documentWithPath:documentID];
   [self writeDocumentRef:docRef2 data:@{@"foo" : @"bar"}];
+  [self awaitExpectation:seen];
 
   [registration remove];
 }

Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm

@@ -311,7 +311,7 @@ - (void)primeBackend {
         }];
 
     // Wait for watch to initialize and deliver first event.
-    [self awaitExpectations];
+    [self awaitExpectation:watchInitialized];
 
     watchUpdateReceived = [self expectationWithDescription:@"Prime backend: Watch update received"];
 
@@ -340,17 +340,18 @@ - (void)primeBackend {
 }
 
 - (void)terminateFirestore:(FIRFirestore *)firestore {
-  [firestore terminateWithCompletion:[self completionForExpectationWithName:@"shutdown"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"shutdown"];
+  [firestore terminateWithCompletion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (void)deleteApp:(FIRApp *)app {
-  XCTestExpectation *expectation = [self expectationWithDescription:@"Delete app"];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"deleteApp"];
   [app deleteApp:^(BOOL completion) {
     XCTAssertTrue(completion);
     [expectation fulfill];
   }];
-  [self awaitExpectations];
+  [self awaitExpectation:expectation];
 }
 
 - (NSString *)documentPath {
@@ -412,7 +413,7 @@ - (FIRDocumentSnapshot *)readDocumentForRef:(FIRDocumentReference *)ref
                     result = doc;
                     [expectation fulfill];
                   }];
-  [self awaitExpectations];
+  [self awaitExpectation:expectation];
 
   return result;
 }
@@ -431,7 +432,7 @@ - (FIRQuerySnapshot *)readDocumentSetForRef:(FIRQuery *)query source:(FIRFiresto
                        result = documentSet;
                        [expectation fulfill];
                      }];
-  [self awaitExpectations];
+  [self awaitExpectation:expectation];
 
   return result;
 }
@@ -452,61 +453,63 @@ - (FIRDocumentSnapshot *)readSnapshotForRef:(FIRDocumentReference *)ref
                                              }
                                            }];
 
-  [self awaitExpectations];
+  [self awaitExpectation:expectation];
   [listener remove];
 
   return result;
 }
 
 - (void)writeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data {
-  [ref setData:data completion:[self completionForExpectationWithName:@"setData"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"setData"];
+  [ref setData:data completion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (void)updateDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<id, id> *)data {
-  [ref updateData:data completion:[self completionForExpectationWithName:@"updateData"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"updateData"];
+  [ref updateData:data completion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (void)deleteDocumentRef:(FIRDocumentReference *)ref {
-  [ref deleteDocumentWithCompletion:[self completionForExpectationWithName:@"deleteDocument"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"deleteDocument"];
+  [ref deleteDocumentWithCompletion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (FIRDocumentReference *)addDocumentRef:(FIRCollectionReference *)ref
                                     data:(NSDictionary<NSString *, id> *)data {
-  FIRDocumentReference *doc =
-      [ref addDocumentWithData:data
-                    completion:[self completionForExpectationWithName:@"addDocument"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"addDocument"];
+  FIRDocumentReference *doc = [ref addDocumentWithData:data
+                                            completion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
   return doc;
 }
 
 - (void)mergeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data {
-  [ref setData:data
-           merge:YES
-      completion:[self completionForExpectationWithName:@"setDataWithMerge"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"setDataWithMerge"];
+  [ref setData:data merge:YES completion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (void)mergeDocumentRef:(FIRDocumentReference *)ref
                     data:(NSDictionary<NSString *, id> *)data
                   fields:(NSArray<id> *)fields {
-  [ref setData:data
-      mergeFields:fields
-       completion:[self completionForExpectationWithName:@"setDataWithMerge"]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"setDataWithMerge"];
+  [ref setData:data mergeFields:fields completion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (void)disableNetwork {
-  [self.db
-      disableNetworkWithCompletion:[self completionForExpectationWithName:@"Disable Network."]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"disableNetwork"];
+  [self.db disableNetworkWithCompletion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (void)enableNetwork {
-  [self.db enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable Network."]];
-  [self awaitExpectations];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"enableNetwork"];
+  [self.db enableNetworkWithCompletion:[self completionForExpectation:expectation]];
+  [self awaitExpectation:expectation];
 }
 
 - (const std::shared_ptr<util::AsyncQueue> &)queueForFirestore:(FIRFirestore *)firestore {

Firestore/Example/Tests/Util/XCTestCase+Await.h

@@ -38,6 +38,12 @@ typedef void (^FSTVoidErrorBlock)(NSError *_Nullable error);
  */
 - (void)awaitExpectations;
 
+/**
+ * Await a specific expectation with a reasonable timeout. If the expectation fails, XCTFail the
+ * test.
+ */
+- (void)awaitExpectation:(XCTestExpectation *)expectation;
+
 /**
  * Returns a reasonable timeout for testing against Firestore.
  */
@@ -49,6 +55,11 @@ typedef void (^FSTVoidErrorBlock)(NSError *_Nullable error);
  */
 - (FSTVoidErrorBlock)completionForExpectationWithName:(NSString *)expectationName;
 
+/**
+ * Returns a completion block that fulfills the given expectation.
+ */
+- (FSTVoidErrorBlock)completionForExpectation:(XCTestExpectation *)expectation;
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Example/Tests/Util/XCTestCase+Await.mm

@@ -36,12 +36,20 @@ - (void)awaitExpectations {
                                }];
 }
 
+- (void)awaitExpectation:(XCTestExpectation *)expectation {
+  [self waitForExpectations:@[ expectation ] timeout:kExpectationWaitSeconds];
+}
+
 - (double)defaultExpectationWaitSeconds {
   return kExpectationWaitSeconds;
 }
 
 - (FSTVoidErrorBlock)completionForExpectationWithName:(NSString *)expectationName {
   XCTestExpectation *expectation = [self expectationWithDescription:expectationName];
+  return [self completionForExpectation:expectation];
+}
+
+- (FSTVoidErrorBlock)completionForExpectation:(XCTestExpectation *)expectation {
   return ^(NSError *error) {
     XCTAssertNil(error);
     [expectation fulfill];

Firestore/core/src/api/query_core.cc

@@ -163,7 +163,7 @@ std::unique_ptr<ListenerRegistration> Query::AddSnapshotListener(
       QuerySnapshot result(firestore_, query_, std::move(snapshot),
                            std::move(metadata));
 
-      user_listener_->OnEvent(result);
+      user_listener_->OnEvent(std::move(result));
     }
 
    private:

Firestore/core/src/core/event_listener.h

@@ -17,8 +17,8 @@
 #ifndef FIRESTORE_CORE_SRC_CORE_EVENT_LISTENER_H_
 #define FIRESTORE_CORE_SRC_CORE_EVENT_LISTENER_H_
 
-#include <atomic>
 #include <memory>
+#include <mutex>  // NOLINT(build/c++11)
 #include <utility>
 
 #include "Firestore/core/src/util/executor.h"
@@ -63,9 +63,6 @@ class AsyncEventListener
   AsyncEventListener(const std::shared_ptr<util::Executor>& executor,
                      DelegateListener&& delegate)
       : executor_(executor), delegate_(std::move(delegate)) {
-    // std::atomic's constructor is not atomic, so assign after contruction
-    // (since assignment is atomic).
-    muted_ = false;
   }
 
   static std::shared_ptr<AsyncEventListener<T>> Create(
@@ -86,7 +83,19 @@ class AsyncEventListener
   void Mute();
 
  private:
-  std::atomic<bool> muted_;
+  // PORTING NOTE: Android uses a volatile here but that's not enough in C++.
+  //
+  // In C++, the user can call `ListenerRegistration::Remove` (which calls
+  // `Mute`) and then immediately delete the state backing the listener. Using
+  // a mutex here instead of an atomic ensures that `Mute` won't return until
+  // it's safe to delete the state backing a listener. In Java this is safe
+  // because the state backing the listener is garbage collected so it doesn't
+  // matter if the mute is concurrent with a callback.
+  //
+  // Use a recursive mutex instead of `std::mutex` to avoid deadlock in the case
+  // where a user calls `Remove` from within a callback on that listener.
+  std::recursive_mutex mutex_;
+  bool muted_ = false;
   std::shared_ptr<util::Executor> executor_;
   DelegateListener delegate_;
 };
@@ -119,6 +128,7 @@ std::shared_ptr<AsyncEventListener<T>> AsyncEventListener<T>::Create(
 
 template <typename T>
 void AsyncEventListener<T>::Mute() {
+  std::lock_guard<std::recursive_mutex> lock(mutex_);
   muted_ = true;
 }
 
@@ -131,6 +141,7 @@ void AsyncEventListener<T>::OnEvent(util::StatusOr<T> maybe_value) {
   std::shared_ptr<AsyncEventListener<T>> shared_this = this->shared_from_this();
 
   executor_->Execute([shared_this, maybe_value]() {
+    std::lock_guard<std::recursive_mutex> lock(shared_this->mutex_);
     if (!shared_this->muted_) {
       shared_this->delegate_->OnEvent(std::move(maybe_value));
     }