Fix removeLifecycleObserver in AsyncSocket

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: yfeldblum

Differential Revision: D27033014

fbshipit-source-id: df84bf65207a00f3b76d8950ac2c29f477508da1

folly/io/async/AsyncSocket.cpp

@@ -2531,7 +2531,14 @@ void AsyncSocket::addLifecycleObserver(
   if (eventBase_) {
     eventBase_->dcheckIsInEventBaseThread();
   }
-  lifecycleObservers_.push_back(observer);
+
+  // adding the same observer multiple times is not allowed
+  auto& observers = lifecycleObservers_;
+  CHECK(
+      std::find(observers.begin(), observers.end(), observer) ==
+      observers.end());
+
+  observers.push_back(observer);
   observer->observerAttach(this);
   if (observer->getConfig().byteEvents) {
     if (byteEventHelper_ && byteEventHelper_->maybeEx.has_value()) {
@@ -2547,16 +2554,13 @@ void AsyncSocket::addLifecycleObserver(
 
 bool AsyncSocket::removeLifecycleObserver(
     AsyncTransport::LifecycleObserver* observer) {
-  const auto eraseIt = std::remove(
-      lifecycleObservers_.begin(), lifecycleObservers_.end(), observer);
-  if (eraseIt == lifecycleObservers_.end()) {
+  auto& observers = lifecycleObservers_;
+  auto it = std::find(observers.begin(), observers.end(), observer);
+  if (it == observers.end()) {
     return false;
   }
-
-  for (auto it = eraseIt; it != lifecycleObservers_.end(); it++) {
-    (*it)->observerDetach(this);
-  }
-  lifecycleObservers_.erase(eraseIt, lifecycleObservers_.end());
+  observer->observerDetach(this);
+  observers.erase(it);
   return true;
 }
 

folly/io/async/test/AsyncSocketTest2.cpp

@@ -5351,9 +5351,10 @@ TEST(AsyncSocket, LifecycleObserverMultipleAttachThenDestroySocket) {
 }
 
 TEST(AsyncSocket, LifecycleObserverAttachRemove) {
-  auto cb = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
   EventBase evb;
   auto socket = AsyncSocket::UniquePtr(new AsyncSocket(&evb));
+
+  auto cb = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
   EXPECT_CALL(*cb, observerAttachMock(socket.get()));
   socket->addLifecycleObserver(cb.get());
   Mock::VerifyAndClearExpectations(cb.get());
@@ -5365,6 +5366,64 @@ TEST(AsyncSocket, LifecycleObserverAttachRemove) {
   Mock::VerifyAndClearExpectations(cb.get());
 }
 
+TEST(AsyncSocket, LifecycleObserverAttachRemoveMultiple) {
+  EventBase evb;
+  auto socket = AsyncSocket::UniquePtr(new AsyncSocket(&evb));
+
+  auto cb1 = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
+  EXPECT_CALL(*cb1, observerAttachMock(socket.get()));
+  socket->addLifecycleObserver(cb1.get());
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(socket->getLifecycleObservers(), UnorderedElementsAre(cb1.get()));
+
+  auto cb2 = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
+  EXPECT_CALL(*cb2, observerAttachMock(socket.get()));
+  socket->addLifecycleObserver(cb2.get());
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(
+      socket->getLifecycleObservers(),
+      UnorderedElementsAre(cb1.get(), cb2.get()));
+
+  EXPECT_CALL(*cb1, observerDetachMock(socket.get()));
+  EXPECT_TRUE(socket->removeLifecycleObserver(cb1.get()));
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(socket->getLifecycleObservers(), UnorderedElementsAre(cb2.get()));
+
+  EXPECT_CALL(*cb2, observerDetachMock(socket.get()));
+  EXPECT_TRUE(socket->removeLifecycleObserver(cb2.get()));
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(socket->getLifecycleObservers(), IsEmpty());
+}
+
+TEST(AsyncSocket, LifecycleObserverAttachRemoveMultipleReverse) {
+  EventBase evb;
+  auto socket = AsyncSocket::UniquePtr(new AsyncSocket(&evb));
+
+  auto cb1 = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
+  EXPECT_CALL(*cb1, observerAttachMock(socket.get()));
+  socket->addLifecycleObserver(cb1.get());
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(socket->getLifecycleObservers(), UnorderedElementsAre(cb1.get()));
+
+  auto cb2 = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
+  EXPECT_CALL(*cb2, observerAttachMock(socket.get()));
+  socket->addLifecycleObserver(cb2.get());
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(
+      socket->getLifecycleObservers(),
+      UnorderedElementsAre(cb1.get(), cb2.get()));
+
+  EXPECT_CALL(*cb2, observerDetachMock(socket.get()));
+  EXPECT_TRUE(socket->removeLifecycleObserver(cb2.get()));
+  Mock::VerifyAndClearExpectations(cb2.get());
+  EXPECT_THAT(socket->getLifecycleObservers(), UnorderedElementsAre(cb1.get()));
+
+  EXPECT_CALL(*cb1, observerDetachMock(socket.get()));
+  EXPECT_TRUE(socket->removeLifecycleObserver(cb1.get()));
+  Mock::VerifyAndClearExpectations(cb1.get());
+  EXPECT_THAT(socket->getLifecycleObservers(), IsEmpty());
+}
+
 TEST(AsyncSocket, LifecycleObserverRemoveMissing) {
   auto cb = std::make_unique<StrictMock<MockAsyncSocketLifecycleObserver>>();
   EventBase evb;