Fix removeObserver and removeAcceptObserver

Summary: When multiple observers attached and one is removed, removal code can end up calling `observerDetach` on the wrong observer. Simplified and fixed removal logic and added new tests.

Reviewed By: yangchi

Differential Revision: D27033221

fbshipit-source-id: d200fd2243a678890758b2652b61d16887f073dd

quic/api/QuicTransportBase.cpp

@@ -2627,21 +2627,22 @@ void QuicTransportBase::resetNonControlStreams(
 }
 
 void QuicTransportBase::addObserver(Observer* observer) {
+  // adding the same observer multiple times is not allowed
+  CHECK(
+      std::find(observers_->begin(), observers_->end(), observer) ==
+      observers_->end());
+
   observers_->push_back(CHECK_NOTNULL(observer));
   observer->observerAttach(this);
 }
 
 bool QuicTransportBase::removeObserver(Observer* observer) {
-  const auto eraseIt =
+  auto it = std::find(observers_->begin(), observers_->end(), observer);
-      std::remove(observers_->begin(), observers_->end(), observer);
+  if (it == observers_->end()) {
-  if (eraseIt == observers_->end()) {
     return false;
   }
-
+  observer->observerDetach(this);
-  for (auto it = eraseIt; it != observers_->end(); it++) {
+  observers_->erase(it);
-    (*it)->observerDetach(this);
-  }
-  observers_->erase(eraseIt, observers_->end());
   return true;
 }
 

quic/api/test/QuicTransportBaseTest.cpp

@@ -3497,6 +3497,56 @@ TEST_F(QuicTransportImplTest, ObserverAttachRemove) {
   EXPECT_THAT(transport->getObservers(), IsEmpty());
 }
 
+TEST_F(QuicTransportImplTest, ObserverAttachRemoveMultiple) {
+  auto cb1 = std::make_unique<StrictMock<MockObserver>>();
+  EXPECT_CALL(*cb1, observerAttach(transport.get()));
+  transport->addObserver(cb1.get());
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(transport->getObservers(), UnorderedElementsAre(cb1.get()));
+
+  auto cb2 = std::make_unique<StrictMock<MockObserver>>();
+  EXPECT_CALL(*cb2, observerAttach(transport.get()));
+  transport->addObserver(cb2.get());
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(
+      transport->getObservers(), UnorderedElementsAre(cb1.get(), cb2.get()));
+
+  EXPECT_CALL(*cb1, observerDetach(transport.get()));
+  EXPECT_TRUE(transport->removeObserver(cb1.get()));
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(transport->getObservers(), UnorderedElementsAre(cb2.get()));
+
+  EXPECT_CALL(*cb2, observerDetach(transport.get()));
+  EXPECT_TRUE(transport->removeObserver(cb2.get()));
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(transport->getObservers(), IsEmpty());
+}
+
+TEST_F(QuicTransportImplTest, ObserverAttachRemoveMultipleReverse) {
+  auto cb1 = std::make_unique<StrictMock<MockObserver>>();
+  EXPECT_CALL(*cb1, observerAttach(transport.get()));
+  transport->addObserver(cb1.get());
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(transport->getObservers(), UnorderedElementsAre(cb1.get()));
+
+  auto cb2 = std::make_unique<StrictMock<MockObserver>>();
+  EXPECT_CALL(*cb2, observerAttach(transport.get()));
+  transport->addObserver(cb2.get());
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(
+      transport->getObservers(), UnorderedElementsAre(cb1.get(), cb2.get()));
+
+  EXPECT_CALL(*cb2, observerDetach(transport.get()));
+  EXPECT_TRUE(transport->removeObserver(cb2.get()));
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(transport->getObservers(), UnorderedElementsAre(cb1.get()));
+
+  EXPECT_CALL(*cb1, observerDetach(transport.get()));
+  EXPECT_TRUE(transport->removeObserver(cb1.get()));
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(transport->getObservers(), IsEmpty());
+}
+
 TEST_F(QuicTransportImplTest, ObserverRemoveMissing) {
   auto cb = std::make_unique<StrictMock<MockObserver>>();
   EXPECT_FALSE(transport->removeObserver(cb.get()));

quic/server/QuicServerWorker.cpp

@@ -1289,21 +1289,22 @@ QuicServerWorker::AcceptObserverList::~AcceptObserverList() {
 }
 
 void QuicServerWorker::AcceptObserverList::add(AcceptObserver* observer) {
+  // adding the same observer multiple times is not allowed
+  CHECK(
+      std::find(observers_.begin(), observers_.end(), observer) ==
+      observers_.end());
+
   observers_.emplace_back(CHECK_NOTNULL(observer));
   observer->observerAttach(worker_);
 }
 
 bool QuicServerWorker::AcceptObserverList::remove(AcceptObserver* observer) {
-  const auto eraseIt =
+  auto it = std::find(observers_.begin(), observers_.end(), observer);
-      std::remove(observers_.begin(), observers_.end(), observer);
+  if (it == observers_.end()) {
-  if (eraseIt == observers_.end()) {
     return false;
   }
-
+  observer->observerDetach(worker_);
-  for (auto it = eraseIt; it != observers_.end(); it++) {
+  observers_.erase(it);
-    (*it)->observerDetach(worker_);
-  }
-  observers_.erase(eraseIt, observers_.end());
   return true;
 }
 

quic/server/test/QuicServerTest.cpp

@@ -1186,20 +1186,37 @@ TEST_F(QuicServerWorkerTest, AcceptObserverMultipleRemove) {
   auto cb2 = std::make_unique<StrictMock<MockAcceptObserver>>();
   EXPECT_CALL(*cb2, observerAttach(worker_.get()));
   worker_->addAcceptObserver(cb2.get());
-  Mock::VerifyAndClearExpectations(cb1.get());
-  Mock::VerifyAndClearExpectations(cb2.get());
-
-  EXPECT_CALL(*cb2, observerDetach(worker_.get()));
-  EXPECT_TRUE(worker_->removeAcceptObserver(cb2.get()));
-  Mock::VerifyAndClearExpectations(cb1.get());
   Mock::VerifyAndClearExpectations(cb2.get());
 
   EXPECT_CALL(*cb1, observerDetach(worker_.get()));
   EXPECT_TRUE(worker_->removeAcceptObserver(cb1.get()));
   Mock::VerifyAndClearExpectations(cb1.get());
+
+  EXPECT_CALL(*cb2, observerDetach(worker_.get()));
+  EXPECT_TRUE(worker_->removeAcceptObserver(cb2.get()));
   Mock::VerifyAndClearExpectations(cb2.get());
 }
 
+TEST_F(QuicServerWorkerTest, AcceptObserverMultipleRemoveReverse) {
+  auto cb1 = std::make_unique<StrictMock<MockAcceptObserver>>();
+  EXPECT_CALL(*cb1, observerAttach(worker_.get()));
+  worker_->addAcceptObserver(cb1.get());
+  Mock::VerifyAndClearExpectations(cb1.get());
+
+  auto cb2 = std::make_unique<StrictMock<MockAcceptObserver>>();
+  EXPECT_CALL(*cb2, observerAttach(worker_.get()));
+  worker_->addAcceptObserver(cb2.get());
+  Mock::VerifyAndClearExpectations(cb2.get());
+
+  EXPECT_CALL(*cb2, observerDetach(worker_.get()));
+  EXPECT_TRUE(worker_->removeAcceptObserver(cb2.get()));
+  Mock::VerifyAndClearExpectations(cb2.get());
+
+  EXPECT_CALL(*cb1, observerDetach(worker_.get()));
+  EXPECT_TRUE(worker_->removeAcceptObserver(cb1.get()));
+  Mock::VerifyAndClearExpectations(cb1.get());
+}
+
 TEST_F(QuicServerWorkerTest, AcceptObserverMultipleAcceptorDestroyed) {
   auto cb1 = std::make_unique<StrictMock<MockAcceptObserver>>();
   EXPECT_CALL(*cb1, observerAttach(worker_.get()));
@@ -1209,7 +1226,6 @@ TEST_F(QuicServerWorkerTest, AcceptObserverMultipleAcceptorDestroyed) {
   auto cb2 = std::make_unique<StrictMock<MockAcceptObserver>>();
   EXPECT_CALL(*cb2, observerAttach(worker_.get()));
   worker_->addAcceptObserver(cb2.get());
-  Mock::VerifyAndClearExpectations(cb1.get());
   Mock::VerifyAndClearExpectations(cb2.get());
 
   // destroy the acceptor while the AcceptObserver is installed