Avoid leaking connection if connection.close is received during shutdown

src/rmq/rmqamqp/rmqamqp_connection.cpp

@@ -242,6 +242,7 @@ Connection::Connection(
 , d_timerFactory(timerFactory)
 , d_firstConnectCb()
 , d_closeCb()
+, d_closeRequested(false)
 , d_connectStartTime()
 , d_hasBeenConnected(false)
 , d_vhostTags()
@@ -255,15 +256,24 @@ void Connection::close(const CloseFinishCallback& closeCallback)
 {
     BALL_LOG_INFO << "Closing connection: " << connectionDebugName();
 
-    d_closeCb = closeCallback;
+    d_closeCb        = closeCallback;
+    d_closeRequested = true;
 
-    if (d_state == DISCONNECTED) {
-        socketShutdown(Connection::FATAL,
-                       rmqio::Connection::GRACEFUL_DISCONNECT);
-    }
-    else {
-        sendConnectionClose(rmqamqpt::Constants::REPLY_SUCCESS,
-                            "Closing Connection");
+    switch (d_state) {
+        case DISCONNECTED:
+            socketShutdown(Connection::FATAL,
+                           rmqio::Connection::GRACEFUL_DISCONNECT);
+            break;
+        case CONNECTION_CLOSE_SENT_EXIT:
+        case CONNECTION_CLOSE_SENT_RECONNECT:
+        case CONNECTION_CLOSE_RECEIVED:
+        case CONNECTION_CLOSEOK_SENT:
+        case CONNECTION_CLOSED:
+            break;
+        default:
+            sendConnectionClose(rmqamqpt::Constants::REPLY_SUCCESS,
+                                "Closing Connection");
+            break;
     }
 }
 
@@ -579,7 +589,7 @@ void Connection::connectionException(
 void Connection::retry(const bsl::weak_ptr<Connection>& weakSelf)
 {
     bsl::shared_ptr<Connection> self = weakSelf.lock();
-    if (!self) {
+    if (!self || self->d_closeRequested) {
         BALL_LOG_TRACE << "Retry handler called during shutdown";
         return;
     }
@@ -609,11 +619,11 @@ void Connection::closeSocket(DisconnectType disconnectType)
 void Connection::socketShutdown(DisconnectType dcType,
                                 rmqio::Connection::ReturnCode code)
 {
+    const bool willRetry = dcType == WILL_RETRY && !d_closeRequested;
+
     if (code != rmqio::Connection::CONNECT_ERROR) {
         BALL_LOG_INFO << "Disconnected from Broker: " << connectionDebugName()
-                      << (dcType == rmqamqp::Connection::WILL_RETRY
-                              ? " Retrying"
-                              : " Connection closed")
+                      << (willRetry ? " Retrying" : " Connection closed")
                       << ". Error code " << code;
     }
 
@@ -633,7 +643,7 @@ void Connection::socketShutdown(DisconnectType dcType,
 
     d_state = Connection::DISCONNECTED;
     BALL_LOG_TRACE << "State now set to: " << d_state;
-    if (dcType == WILL_RETRY) {
+    if (willRetry) {
         // The connection is down and will be retried: either an establishment
         // attempt failed or an established connection was lost. Schedule the
         // retry before invoking the observational hook: the hook is a

src/rmq/rmqamqp/rmqamqp_connection.h

@@ -208,6 +208,7 @@ class Connection : public bsl::enable_shared_from_this<Connection>,
 
     ConnectedCallback d_firstConnectCb;
     CloseFinishCallback d_closeCb;
+    bool d_closeRequested;
 
     bsls::TimeInterval d_connectStartTime;
     bool d_hasBeenConnected; // For metrics purposes -- indicates whether this

src/tests/rmqamqp/rmqamqp_connection.t.cpp

@@ -993,6 +993,58 @@ TEST_F(ConnectionTests, GracefulClose)
     EXPECT_TRUE(invoked);
 }
 
+TEST_F(ConnectionTests, NoReconnectWhenServerClosesDuringGracefulClose)
+{
+    expectFirstHandshakeFrames();
+
+    bsl::shared_ptr<rmqamqp::Connection> conn = createAndStartConnection();
+
+    d_eventLoop.run();
+    d_eventLoop.restart();
+
+    expectCloseFrame(false);
+    expectCloseFrame();
+    expectCloseOkFrame();
+
+    EXPECT_CALL(*d_retryHandler, retry(_)).Times(0);
+
+    bool invoked = false;
+    conn->close(bdlf::BindUtil::bind(&gracefulCloseHandler, bsl::ref(invoked)));
+
+    d_eventLoop.run();
+
+    EXPECT_THAT(conn->state(), Eq(rmqamqp::Connection::DISCONNECTED));
+    EXPECT_THAT(d_replayFrame.getLength(), Eq(0));
+    EXPECT_TRUE(invoked);
+}
+
+TEST_F(ConnectionTests, NoCloseWhenServerAlreadyClosed)
+{
+    expectFirstHandshakeFrames();
+
+    bsl::shared_ptr<rmqamqp::Connection> conn = createAndStartConnection();
+
+    d_eventLoop.run();
+    d_eventLoop.restart();
+
+    expectCloseFrame();
+    expectCloseOkFrame();
+
+    EXPECT_CALL(*d_retryHandler, retry(_)).Times(0);
+
+    feedNextFrame();
+    d_eventLoop.run_one();
+
+    bool invoked = false;
+    conn->close(bdlf::BindUtil::bind(&gracefulCloseHandler, bsl::ref(invoked)));
+
+    d_eventLoop.run();
+
+    EXPECT_THAT(conn->state(), Eq(rmqamqp::Connection::DISCONNECTED));
+    EXPECT_THAT(d_replayFrame.getLength(), Eq(0));
+    EXPECT_TRUE(invoked);
+}
+
 TEST_F(ConnectionTests, AddReceiveChannelWhenConnected)
 {
     expectFirstHandshakeFrames();