Skip to content

Commit 59c8c95

Browse files
wilhuffwilhuff
authored
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.
1 parent 60293c6 commit 59c8c95

File tree

7 files changed

+170
-37
lines changed

7 files changed

+170
-37
lines changed

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

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -566,6 +566,7 @@ - (void)testSnapshotsInSyncListenerFiresAfterListenersInSync {
566566
}];
567567

568568
[self writeDocumentRef:ref data:@{@"foo" : @3}];
569+
[self awaitExpectation:done];
569570
}
570571

571572
- (void)testSnapshotsInSyncRemoveIsIdempotent {
@@ -1484,6 +1485,104 @@ - (void)testCanRemoveListenerAfterTermination {
14841485
[listenerRegistration remove];
14851486
}
14861487

1488+
- (void)testListenerCallbackBlocksRemove {
1489+
// This tests a guarantee required for C++ that doesn't strictly matter for Objective-C and has no
1490+
// equivalent on other platforms.
1491+
//
1492+
// The problem for C++ is that users can register a listener that refers to some state, then call
1493+
// `ListenerRegistration::Remove()` and expect to be able to immediately delete that state. The
1494+
// trouble is that there may be a callback in progress against that listener so the implementation
1495+
// now blocks the remove call until the callback is complete.
1496+
XCTestExpectation *running = [self expectationWithDescription:@"listener running"];
1497+
XCTestExpectation *allowCompletion =
1498+
[self expectationWithDescription:@"allow listener to complete"];
1499+
XCTestExpectation *removing = [self expectationWithDescription:@"attempting to remove listener"];
1500+
XCTestExpectation *removed = [self expectationWithDescription:@"listener removed"];
1501+
1502+
NSMutableString *steps = [NSMutableString string];
1503+
1504+
FIRDocumentReference *doc = [self documentRef];
1505+
[self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
1506+
1507+
__block bool firstTime = true;
1508+
1509+
id<FIRListenerRegistration> listener =
1510+
[doc addSnapshotListener:^(FIRDocumentSnapshot *, NSError *) {
1511+
@synchronized(self) {
1512+
if (!firstTime) {
1513+
return;
1514+
}
1515+
firstTime = false;
1516+
}
1517+
1518+
[steps appendString:@"1"];
1519+
[running fulfill];
1520+
1521+
[self awaitExpectation:allowCompletion];
1522+
[steps appendString:@"3"];
1523+
}];
1524+
1525+
// Call remove asynchronously to avoid blocking the main test thread.
1526+
dispatch_queue_t async = dispatch_queue_create("firestore.async", DISPATCH_QUEUE_SERIAL);
1527+
dispatch_async(async, ^{
1528+
[self awaitExpectation:running];
1529+
[steps appendString:@"2"];
1530+
1531+
[removing fulfill];
1532+
[listener remove];
1533+
1534+
[steps appendString:@"4"];
1535+
[removed fulfill];
1536+
});
1537+
1538+
// Perform a write to `doc` which will trigger the listener callback. Don't wait for completion
1539+
// though because that completion handler is in line behind the listener callback that the test
1540+
// is blocking.
1541+
XCTestExpectation *setData = [self expectationWithDescription:@"setData"];
1542+
[doc setData:@{@"foo" : @"bar"} completion:[self completionForExpectation:setData]];
1543+
1544+
[self awaitExpectation:removing];
1545+
[allowCompletion fulfill];
1546+
1547+
[self awaitExpectation:removed];
1548+
XCTAssertEqualObjects(steps, @"1234");
1549+
1550+
[self awaitExpectation:setData];
1551+
}
1552+
1553+
- (void)testListenerCallbackCanCallRemoveWithoutBlocking {
1554+
// This tests a guarantee required for C++ that doesn't strictly matter for Objective-C and has no
1555+
// equivalent on other platforms. See `testListenerCallbackBlocksRemove` for background.
1556+
XCTestExpectation *removed = [self expectationWithDescription:@"listener removed"];
1557+
1558+
NSMutableString *steps = [NSMutableString string];
1559+
1560+
FIRDocumentReference *doc = [self documentRef];
1561+
[self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
1562+
1563+
__block id<FIRListenerRegistration> listener = nil;
1564+
1565+
@synchronized(self) {
1566+
listener = [doc addSnapshotListener:^(FIRDocumentSnapshot *, NSError *) {
1567+
[steps appendString:@"1"];
1568+
1569+
@synchronized(self) {
1570+
// This test is successful if this method does not block.
1571+
[listener remove];
1572+
}
1573+
1574+
[steps appendString:@"2"];
1575+
[removed fulfill];
1576+
}];
1577+
}
1578+
1579+
// Perform a write to `doc` which will trigger the listener callback.
1580+
[self writeDocumentRef:doc data:@{@"foo" : @"bar2"}];
1581+
1582+
[self awaitExpectation:removed];
1583+
XCTAssertEqualObjects(steps, @"12");
1584+
}
1585+
14871586
- (void)testWaitForPendingWritesCompletes {
14881587
FIRDocumentReference *doc = [self documentRef];
14891588
FIRFirestore *firestore = doc.firestore;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,7 @@ - (void)testCanOutliveDocumentReference {
158158

159159
FIRDocumentReference *docRef2 = [collectionRef documentWithPath:documentID];
160160
[self writeDocumentRef:docRef2 data:@{@"foo" : @"bar"}];
161+
[self awaitExpectation:seen];
161162

162163
[registration remove];
163164
}

‎Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm

Lines changed: 34 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ - (void)primeBackend {
311311
}];
312312

313313
// Wait for watch to initialize and deliver first event.
314-
[self awaitExpectations];
314+
[self awaitExpectation:watchInitialized];
315315

316316
watchUpdateReceived = [self expectationWithDescription:@"Prime backend: Watch update received"];
317317

@@ -340,17 +340,18 @@ - (void)primeBackend {
340340
}
341341

342342
- (void)terminateFirestore:(FIRFirestore *)firestore {
343-
[firestore terminateWithCompletion:[self completionForExpectationWithName:@"shutdown"]];
344-
[self awaitExpectations];
343+
XCTestExpectation *expectation = [self expectationWithDescription:@"shutdown"];
344+
[firestore terminateWithCompletion:[self completionForExpectation:expectation]];
345+
[self awaitExpectation:expectation];
345346
}
346347

347348
- (void)deleteApp:(FIRApp *)app {
348-
XCTestExpectation *expectation = [self expectationWithDescription:@"Delete app"];
349+
XCTestExpectation *expectation = [self expectationWithDescription:@"deleteApp"];
349350
[app deleteApp:^(BOOL completion) {
350351
XCTAssertTrue(completion);
351352
[expectation fulfill];
352353
}];
353-
[self awaitExpectations];
354+
[self awaitExpectation:expectation];
354355
}
355356

356357
- (NSString *)documentPath {
@@ -412,7 +413,7 @@ - (FIRDocumentSnapshot *)readDocumentForRef:(FIRDocumentReference *)ref
412413
result = doc;
413414
[expectation fulfill];
414415
}];
415-
[self awaitExpectations];
416+
[self awaitExpectation:expectation];
416417

417418
return result;
418419
}
@@ -431,7 +432,7 @@ - (FIRQuerySnapshot *)readDocumentSetForRef:(FIRQuery *)query source:(FIRFiresto
431432
result = documentSet;
432433
[expectation fulfill];
433434
}];
434-
[self awaitExpectations];
435+
[self awaitExpectation:expectation];
435436

436437
return result;
437438
}
@@ -452,61 +453,63 @@ - (FIRDocumentSnapshot *)readSnapshotForRef:(FIRDocumentReference *)ref
452453
}
453454
}];
454455

455-
[self awaitExpectations];
456+
[self awaitExpectation:expectation];
456457
[listener remove];
457458

458459
return result;
459460
}
460461

461462
- (void)writeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data {
462-
[ref setData:data completion:[self completionForExpectationWithName:@"setData"]];
463-
[self awaitExpectations];
463+
XCTestExpectation *expectation = [self expectationWithDescription:@"setData"];
464+
[ref setData:data completion:[self completionForExpectation:expectation]];
465+
[self awaitExpectation:expectation];
464466
}
465467

466468
- (void)updateDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<id, id> *)data {
467-
[ref updateData:data completion:[self completionForExpectationWithName:@"updateData"]];
468-
[self awaitExpectations];
469+
XCTestExpectation *expectation = [self expectationWithDescription:@"updateData"];
470+
[ref updateData:data completion:[self completionForExpectation:expectation]];
471+
[self awaitExpectation:expectation];
469472
}
470473

471474
- (void)deleteDocumentRef:(FIRDocumentReference *)ref {
472-
[ref deleteDocumentWithCompletion:[self completionForExpectationWithName:@"deleteDocument"]];
473-
[self awaitExpectations];
475+
XCTestExpectation *expectation = [self expectationWithDescription:@"deleteDocument"];
476+
[ref deleteDocumentWithCompletion:[self completionForExpectation:expectation]];
477+
[self awaitExpectation:expectation];
474478
}
475479

476480
- (FIRDocumentReference *)addDocumentRef:(FIRCollectionReference *)ref
477481
data:(NSDictionary<NSString *, id> *)data {
478-
FIRDocumentReference *doc =
479-
[ref addDocumentWithData:data
480-
completion:[self completionForExpectationWithName:@"addDocument"]];
481-
[self awaitExpectations];
482+
XCTestExpectation *expectation = [self expectationWithDescription:@"addDocument"];
483+
FIRDocumentReference *doc = [ref addDocumentWithData:data
484+
completion:[self completionForExpectation:expectation]];
485+
[self awaitExpectation:expectation];
482486
return doc;
483487
}
484488

485489
- (void)mergeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data {
486-
[ref setData:data
487-
merge:YES
488-
completion:[self completionForExpectationWithName:@"setDataWithMerge"]];
489-
[self awaitExpectations];
490+
XCTestExpectation *expectation = [self expectationWithDescription:@"setDataWithMerge"];
491+
[ref setData:data merge:YES completion:[self completionForExpectation:expectation]];
492+
[self awaitExpectation:expectation];
490493
}
491494

492495
- (void)mergeDocumentRef:(FIRDocumentReference *)ref
493496
data:(NSDictionary<NSString *, id> *)data
494497
fields:(NSArray<id> *)fields {
495-
[ref setData:data
496-
mergeFields:fields
497-
completion:[self completionForExpectationWithName:@"setDataWithMerge"]];
498-
[self awaitExpectations];
498+
XCTestExpectation *expectation = [self expectationWithDescription:@"setDataWithMerge"];
499+
[ref setData:data mergeFields:fields completion:[self completionForExpectation:expectation]];
500+
[self awaitExpectation:expectation];
499501
}
500502

501503
- (void)disableNetwork {
502-
[self.db
503-
disableNetworkWithCompletion:[self completionForExpectationWithName:@"Disable Network."]];
504-
[self awaitExpectations];
504+
XCTestExpectation *expectation = [self expectationWithDescription:@"disableNetwork"];
505+
[self.db disableNetworkWithCompletion:[self completionForExpectation:expectation]];
506+
[self awaitExpectation:expectation];
505507
}
506508

507509
- (void)enableNetwork {
508-
[self.db enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable Network."]];
509-
[self awaitExpectations];
510+
XCTestExpectation *expectation = [self expectationWithDescription:@"enableNetwork"];
511+
[self.db enableNetworkWithCompletion:[self completionForExpectation:expectation]];
512+
[self awaitExpectation:expectation];
510513
}
511514

512515
- (const std::shared_ptr<util::AsyncQueue> &)queueForFirestore:(FIRFirestore *)firestore {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,12 @@ typedef void (^FSTVoidErrorBlock)(NSError *_Nullable error);
3838
*/
3939
- (void)awaitExpectations;
4040

41+
/**
42+
* Await a specific expectation with a reasonable timeout. If the expectation fails, XCTFail the
43+
* test.
44+
*/
45+
- (void)awaitExpectation:(XCTestExpectation *)expectation;
46+
4147
/**
4248
* Returns a reasonable timeout for testing against Firestore.
4349
*/
@@ -49,6 +55,11 @@ typedef void (^FSTVoidErrorBlock)(NSError *_Nullable error);
4955
*/
5056
- (FSTVoidErrorBlock)completionForExpectationWithName:(NSString *)expectationName;
5157

58+
/**
59+
* Returns a completion block that fulfills the given expectation.
60+
*/
61+
- (FSTVoidErrorBlock)completionForExpectation:(XCTestExpectation *)expectation;
62+
5263
@end
5364

5465
NS_ASSUME_NONNULL_END

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,12 +36,20 @@ - (void)awaitExpectations {
3636
}];
3737
}
3838

39+
- (void)awaitExpectation:(XCTestExpectation *)expectation {
40+
[self waitForExpectations:@[ expectation ] timeout:kExpectationWaitSeconds];
41+
}
42+
3943
- (double)defaultExpectationWaitSeconds {
4044
return kExpectationWaitSeconds;
4145
}
4246

4347
- (FSTVoidErrorBlock)completionForExpectationWithName:(NSString *)expectationName {
4448
XCTestExpectation *expectation = [self expectationWithDescription:expectationName];
49+
return [self completionForExpectation:expectation];
50+
}
51+
52+
- (FSTVoidErrorBlock)completionForExpectation:(XCTestExpectation *)expectation {
4553
return ^(NSError *error) {
4654
XCTAssertNil(error);
4755
[expectation fulfill];

‎Firestore/core/src/api/query_core.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@ std::unique_ptr<ListenerRegistration> Query::AddSnapshotListener(
163163
QuerySnapshot result(firestore_, query_, std::move(snapshot),
164164
std::move(metadata));
165165

166-
user_listener_->OnEvent(result);
166+
user_listener_->OnEvent(std::move(result));
167167
}
168168

169169
private:

‎Firestore/core/src/core/event_listener.h

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@
1717
#ifndef FIRESTORE_CORE_SRC_CORE_EVENT_LISTENER_H_
1818
#define FIRESTORE_CORE_SRC_CORE_EVENT_LISTENER_H_
1919

20-
#include <atomic>
2120
#include <memory>
21+
#include <mutex> // NOLINT(build/c++11)
2222
#include <utility>
2323

2424
#include "Firestore/core/src/util/executor.h"
@@ -63,9 +63,6 @@ class AsyncEventListener
6363
AsyncEventListener(const std::shared_ptr<util::Executor>& executor,
6464
DelegateListener&& delegate)
6565
: executor_(executor), delegate_(std::move(delegate)) {
66-
// std::atomic's constructor is not atomic, so assign after contruction
67-
// (since assignment is atomic).
68-
muted_ = false;
6966
}
7067

7168
static std::shared_ptr<AsyncEventListener<T>> Create(
@@ -86,7 +83,19 @@ class AsyncEventListener
8683
void Mute();
8784

8885
private:
89-
std::atomic<bool> muted_;
86+
// PORTING NOTE: Android uses a volatile here but that's not enough in C++.
87+
//
88+
// In C++, the user can call `ListenerRegistration::Remove` (which calls
89+
// `Mute`) and then immediately delete the state backing the listener. Using
90+
// a mutex here instead of an atomic ensures that `Mute` won't return until
91+
// it's safe to delete the state backing a listener. In Java this is safe
92+
// because the state backing the listener is garbage collected so it doesn't
93+
// matter if the mute is concurrent with a callback.
94+
//
95+
// Use a recursive mutex instead of `std::mutex` to avoid deadlock in the case
96+
// where a user calls `Remove` from within a callback on that listener.
97+
std::recursive_mutex mutex_;
98+
bool muted_ = false;
9099
std::shared_ptr<util::Executor> executor_;
91100
DelegateListener delegate_;
92101
};
@@ -119,6 +128,7 @@ std::shared_ptr<AsyncEventListener<T>> AsyncEventListener<T>::Create(
119128

120129
template <typename T>
121130
void AsyncEventListener<T>::Mute() {
131+
std::lock_guard<std::recursive_mutex> lock(mutex_);
122132
muted_ = true;
123133
}
124134

@@ -131,6 +141,7 @@ void AsyncEventListener<T>::OnEvent(util::StatusOr<T> maybe_value) {
131141
std::shared_ptr<AsyncEventListener<T>> shared_this = this->shared_from_this();
132142

133143
executor_->Execute([shared_this, maybe_value]() {
144+
std::lock_guard<std::recursive_mutex> lock(shared_this->mutex_);
134145
if (!shared_this->muted_) {
135146
shared_this->delegate_->OnEvent(std::move(maybe_value));
136147
}

0 commit comments

Comments
 (0)