wait for pending operations gracefully

Author: BewareMyPower

lib/ClientConnection.cc

@@ -411,7 +411,18 @@ typedef ASIO::detail::socket_option::integer<IPPROTO_TCP, TCP_KEEPIDLE> tcp_keep
  *  at this point the connection is deemed valid to be used by clients of this class
  */
 void ClientConnection::handleTcpConnected(const ASIO_ERROR& err, const tcp::endpoint& endpoint) {
+    if (connectTimeoutTask_) {
+        connectTimeoutTask_->stop();
+    }
+    Lock lock(mutex_);
     if (!err) {
+        if (isClosed()) {
+            ASIO_ERROR err;
+            socket_->close(err);
+            connectPromise_.setFailed(ResultAlreadyClosed);
+            LOG_INFO(cnxString_ << "Connection already closed");
+            return;
+        }
         std::stringstream cnxStringStream;
         try {
             cnxStringStream << "[" << socket_->local_endpoint() << " -> " << socket_->remote_endpoint()
@@ -430,11 +441,6 @@ void ClientConnection::handleTcpConnected(const ASIO_ERROR& err, const tcp::endp
                                 << ", physical address:" << physicalAddress_);
         }
 
-        Lock lock(mutex_);
-        if (isClosed()) {
-            LOG_INFO(cnxString_ << "Connection already closed");
-            return;
-        }
         state_ = TcpConnected;
         lock.unlock();
 
@@ -484,10 +490,12 @@ void ClientConnection::handleTcpConnected(const ASIO_ERROR& err, const tcp::endp
             auto tlsSocket = tlsSocket_;
             // socket and ssl::stream objects must exist until async_handshake is done, otherwise segmentation
             // fault might happen
-            auto callback = [weakSelf, socket, tlsSocket](const ASIO_ERROR& err) {
+            pendingOperations_++;
+            auto callback = [this, weakSelf, socket, tlsSocket](const ASIO_ERROR& err) {
                 auto self = weakSelf.lock();
                 if (self) {
-                    self->handleHandshake(err);
+                    handleHandshake(err);
+                    pendingOperations_--;
                 }
             };
             tlsSocket_->async_handshake(ASIO::ssl::stream<tcp::socket>::client,
@@ -496,6 +504,7 @@ void ClientConnection::handleTcpConnected(const ASIO_ERROR& err, const tcp::endp
             handleHandshake(ASIO_SUCCESS);
         }
     } else {
+        lock.unlock();
         LOG_ERROR(cnxString_ << "Failed to establish connection to " << endpoint << ": " << err.message());
         if (err == ASIO::error::operation_aborted) {
             close();
@@ -596,11 +605,13 @@ void ClientConnection::tcpConnectAsync() {
     LOG_DEBUG(cnxString_ << "Resolving " << service_url.host() << ":" << service_url.port());
 
     auto weakSelf = weak_from_this();
+    pendingOperations_++;
     resolver_->async_resolve(service_url.host(), std::to_string(service_url.port()),
-                             [weakSelf](auto err, const auto& results) {
+                             [this, weakSelf](auto err, const auto& results) {
                                  auto self = weakSelf.lock();
                                  if (self) {
-                                     self->handleResolve(err, results);
+                                     handleResolve(err, results);
+                                     pendingOperations_--;
                                  }
                              });
 }
@@ -642,12 +653,15 @@ void ClientConnection::handleResolve(ASIO_ERROR err, const tcp::resolver::result
         ptr->connectTimeoutTask_->stop();
     });
     connectTimeoutTask_->start();
-    ASIO::async_connect(*socket_, results, [weakSelf](const ASIO_ERROR& err, const tcp::endpoint& endpoint) {
-        auto self = weakSelf.lock();
-        if (self) {
-            self->handleTcpConnected(err, endpoint);
-        }
-    });
+    pendingOperations_++;
+    ASIO::async_connect(*socket_, results,
+                        [this, weakSelf](const ASIO_ERROR& err, const tcp::endpoint& endpoint) {
+                            auto self = weakSelf.lock();
+                            if (self) {
+                                handleTcpConnected(err, endpoint);
+                                pendingOperations_--;
+                            }
+                        });
 }
 
 void ClientConnection::readNextCommand() {
@@ -1079,6 +1093,7 @@ void ClientConnection::sendCommand(const SharedBuffer& cmd) {
 
     if (pendingWriteOperations_++ == 0) {
         // Write immediately to socket
+        lock.unlock();
         if (tlsSocket_) {
             auto weakSelf = weak_from_this();
             auto callback = [weakSelf, cmd]() {
@@ -1162,6 +1177,7 @@ void ClientConnection::sendPendingCommands() {
         pendingWriteBuffers_.pop_front();
 
         auto self = shared_from_this();
+        lock.unlock();
         if (any.type() == typeid(SharedBuffer)) {
             SharedBuffer buffer = std::any_cast<SharedBuffer>(any);
             asyncWrite(buffer.const_asio_buffer(),
@@ -1297,16 +1313,19 @@ void ClientConnection::close(Result result, bool detach) {
     if (isClosed()) {
         return;
     }
-    state_ = Disconnected;
 
     if (socket_) {
         ASIO_ERROR err;
-        socket_->shutdown(ASIO::socket_base::shutdown_both, err);
-        socket_->close(err);
-        if (err) {
-            LOG_WARN(cnxString_ << "Failed to close socket: " << err.message());
-        }
+        if (state_ != Pending) {
+            socket_->shutdown(ASIO::socket_base::shutdown_both, err);
+            socket_->close(err);
+            if (err) {
+                LOG_WARN(cnxString_ << "Failed to close socket: " << err.message());
+            }
+        }  // else: there is an ongoing connect operation, the socket will be closed after the operation
+           // succeeds in `handleTcpConnected`
     }
+    state_ = Disconnected;
     if (tlsSocket_) {
         ASIO_ERROR err;
         tlsSocket_->lowest_layer().close(err);

lib/ClientConnection.h

@@ -221,6 +221,8 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
 
     void handleKeepAliveTimeout();
 
+    auto pendingOperations() const noexcept { return pendingOperations_.load(std::memory_order_acquire); }
+
    private:
     struct PendingRequestData {
         Promise<Result, ResponseData> promise;
@@ -298,25 +300,41 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
 
     template <typename ConstBufferSequence, typename WriteHandler>
     inline void asyncWrite(const ConstBufferSequence& buffers, WriteHandler handler) {
+        Lock lock{mutex_};
         if (isClosed()) {
             return;
         }
+        lock.unlock();
+        pendingOperations_++;
+        auto wrappedHandler = customAllocWriteHandler(
+            [this, handler = std::move(handler)](const ASIO_ERROR& err, size_t bytesTransferred) mutable {
+                handler(err, bytesTransferred);
+                pendingOperations_--;
+            });
         if (tlsSocket_) {
-            ASIO::async_write(*tlsSocket_, buffers, ASIO::bind_executor(strand_, handler));
+            ASIO::async_write(*tlsSocket_, buffers, ASIO::bind_executor(strand_, wrappedHandler));
         } else {
-            ASIO::async_write(*socket_, buffers, handler);
+            ASIO::async_write(*socket_, buffers, wrappedHandler);
         }
     }
 
     template <typename MutableBufferSequence, typename ReadHandler>
     inline void asyncReceive(const MutableBufferSequence& buffers, ReadHandler handler) {
+        Lock lock{mutex_};
         if (isClosed()) {
             return;
         }
+        lock.unlock();
+        pendingOperations_++;
+        auto wrappedHandler = customAllocReadHandler(
+            [this, handler = std::move(handler)](const ASIO_ERROR& err, size_t bytesTransferred) mutable {
+                handler(err, bytesTransferred);
+                pendingOperations_--;
+            });
         if (tlsSocket_) {
-            tlsSocket_->async_read_some(buffers, ASIO::bind_executor(strand_, handler));
+            tlsSocket_->async_read_some(buffers, ASIO::bind_executor(strand_, wrappedHandler));
         } else {
-            socket_->async_receive(buffers, handler);
+            socket_->async_receive(buffers, wrappedHandler);
         }
     }
 
@@ -419,6 +437,7 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     const std::string clientVersion_;
     ConnectionPool& pool_;
     const size_t poolIndex_;
+    std::atomic_int pendingOperations_{0};
 
     friend class PulsarFriend;
     friend class ConsumerTest;

lib/ConnectionPool.cc

@@ -18,6 +18,8 @@
  */
 #include "ConnectionPool.h"
 
+#include <thread>
+
 #ifdef USE_ASIO
 #include <asio/ip/tcp.hpp>
 #include <asio/ssl.hpp>
@@ -61,8 +63,18 @@ bool ConnectionPool::close() {
         if (cnx) {
             // The 2nd argument is false because removing a value during the iteration will cause segfault
             cnx->close(ResultDisconnected, false);
+            for (int i = 0; i < 5000 && cnx->pendingOperations() > 0; i++) {
+                using namespace std::chrono_literals;
+                std::this_thread::sleep_for(1ms);
+            }
+            if (cnx->pendingOperations() > 0) {
+                LOG_WARN(
+                    "Connection still has pending operations after waiting for 5 seconds, pending count: "
+                    << cnx->pendingOperations());
+            }
         }
     }
+
     pool_.clear();
     return true;
 }