Don't send onDisconnect callback before resuming

Summary:
If we have an open connection when a resume comes in we drop the old
connection as expected. However we also send a onDisconnected callback to the
server which makes it think it should pause the session. It pauses everything
only to unpause it moments later as we were in the process of resuming.

Reviewed By: phoad

Differential Revision: D9267026

fbshipit-source-id: 8894291a84da3ab0072ed87c1b783246092fcc17

Author: lexs

rsocket/statemachine/RSocketStateMachine.cpp

@@ -154,8 +154,11 @@ bool RSocketStateMachine::resumeServer(
   const int64_t serverDelta =
       resumeManager_->lastSentPosition() - resumeParams.serverPosition;
 
-  std::runtime_error exn{"Connection being resumed, dropping old connection"};
-  disconnect(std::move(exn));
+  if (frameTransport) {
+    stats_->socketDisconnected();
+  }
+  closeFrameTransport(
+      std::runtime_error{"Connection being resumed, dropping old connection"});
   setProtocolVersionOrThrow(resumeParams.protocolVersion, frameTransport);
   connect(std::move(frameTransport));
 

rsocket/test/statemachine/RSocketStateMachineTest.cpp

@@ -18,6 +18,7 @@
 #include <yarpl/single/SingleSubscriptions.h>
 #include <yarpl/single/Singles.h>
 #include <yarpl/test_utils/Mocks.h>
+#include "rsocket/RSocketConnectionEvents.h"
 #include "rsocket/RSocketResponder.h"
 #include "rsocket/framing/FrameSerializer_v1_0.h"
 #include "rsocket/framing/FrameTransportImpl.h"
@@ -68,6 +69,11 @@ class ResponderMock : public RSocketResponder {
   }
 };
 
+struct ConnectionEventsMock : public RSocketConnectionEvents {
+  MOCK_METHOD1(onDisconnected, void(const folly::exception_wrapper&));
+  MOCK_METHOD0(onStreamsPaused, void());
+};
+
 class RSocketStateMachineTest : public Test {
  public:
   auto createClient(
@@ -96,6 +102,37 @@ class RSocketStateMachineTest : public Test {
     return stateMachine;
   }
 
+  auto createServer(
+      std::unique_ptr<MockDuplexConnection> connection,
+      std::shared_ptr<RSocketResponder> responder,
+      folly::Optional<ResumeIdentificationToken> resumeToken = folly::none,
+      std::shared_ptr<RSocketConnectionEvents> connectionEvents = nullptr) {
+    auto transport =
+        std::make_shared<FrameTransportImpl>(std::move(connection));
+
+    auto stateMachine = std::make_shared<RSocketStateMachine>(
+        std::move(responder),
+        nullptr,
+        RSocketMode::SERVER,
+        nullptr,
+        std::move(connectionEvents),
+        ResumeManager::makeEmpty(),
+        nullptr);
+
+    if (resumeToken) {
+      SetupParameters setupParameters;
+      setupParameters.resumable = true;
+      setupParameters.token = *resumeToken;
+      stateMachine->connectServer(std::move(transport), setupParameters);
+    } else {
+      SetupParameters setupParameters;
+      setupParameters.resumable = false;
+      stateMachine->connectServer(std::move(transport), setupParameters);
+    }
+
+    return stateMachine;
+  }
+
   const std::unordered_map<StreamId, std::shared_ptr<StreamStateMachineBase>>&
   getStreams(RSocketStateMachine& stateMachine) {
     return stateMachine.streams_;
@@ -370,4 +407,25 @@ TEST_F(RSocketStateMachineTest, TransportOnNextClose) {
   rawTransport->onNext(std::move(buf));
 }
 
+TEST_F(RSocketStateMachineTest, ResumeWithCurrentConnection) {
+  auto resumeToken = ResumeIdentificationToken::generateNew();
+
+  auto eventsMock = std::make_shared<ConnectionEventsMock>();
+  auto stateMachine = createServer(
+      std::make_unique<NiceMock<MockDuplexConnection>>(),
+      std::make_shared<RSocketResponder>(),
+      resumeToken,
+      eventsMock);
+
+  EXPECT_CALL(*eventsMock, onDisconnected(_)).Times(0);
+  EXPECT_CALL(*eventsMock, onStreamsPaused()).Times(0);
+
+  ResumeParameters resumeParams{resumeToken, 0, 0, ProtocolVersion::Latest};
+  auto transport = std::make_shared<FrameTransportImpl>(
+      std::make_unique<NiceMock<MockDuplexConnection>>());
+  stateMachine->resumeServer(transport, resumeParams);
+
+  stateMachine->close({}, StreamCompletionSignal::CONNECTION_END);
+}
+
 } // namespace rsocket