Merge branch 'main' into bewaremypower/service-url-provider

Author: BewareMyPower

.github/workflows/ci-pr-validation.yaml

@@ -260,6 +260,15 @@ jobs:
             Pop-Location
           }
 
+      - name: Ensure vcpkg has full history(windows)
+        if: runner.os == 'Windows'
+        shell: pwsh
+        run: |
+          $isShallow = (git -C "${{ env.VCPKG_ROOT }}" rev-parse --is-shallow-repository).Trim()
+          if ($isShallow -eq "true") {
+            git -C "${{ env.VCPKG_ROOT }}" fetch --unshallow
+          }
+
       - name: remove system vcpkg(windows)
         if: runner.os == 'Windows'
         run: rm -rf "$VCPKG_INSTALLATION_ROOT"

lib/ClientConnection.cc

@@ -26,6 +26,8 @@
 #include <any>
 #include <atomic>
 #include <cstdint>
+#include <future>
+#include <optional>
 #ifdef USE_ASIO
 #include <asio/bind_executor.hpp>
 #include <asio/io_context.hpp>
@@ -42,8 +44,10 @@
 #include <deque>
 #include <functional>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "AsioTimer.h"
@@ -157,14 +161,9 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
      * Close the connection.
      *
      * @param result all pending futures will complete with this result
-     * @param detach remove it from the pool if it's true. When false, the connection remains
-     *        associated with the pool but is logically closed; this is currently used when the
-     *        pool itself is being closed or when switching clusters.
      * @param switchCluster whether the close is triggered by cluster switching
-     *
-     * `detach` should only be false when the connection pool is closed.
      */
-    void close(Result result = ResultConnectError, bool detach = true, bool switchCluster = false);
+    const std::future<void>& close(Result result = ResultConnectError, bool switchCluster = false);
 
     bool isClosed() const;
 
@@ -197,7 +196,7 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
 
     const std::string& brokerAddress() const;
 
-    const std::string& cnxString() const;
+    auto cnxString() const { return *std::atomic_load(&cnxStringPtr_); }
 
     int getServerProtocolVersion() const;
 
@@ -223,28 +222,48 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
         mockingRequests_.store(true, std::memory_order_release);
     }
 
-    void handleKeepAliveTimeout();
+    void handleKeepAliveTimeout(const ASIO_ERROR& ec);
 
    private:
     struct PendingRequestData {
         Promise<Result, ResponseData> promise;
         DeadlineTimerPtr timer;
         std::shared_ptr<std::atomic_bool> hasGotResponse{std::make_shared<std::atomic_bool>(false)};
+
+        void fail(Result result) {
+            cancelTimer(*timer);
+            promise.setFailed(result);
+        }
     };
 
     struct LookupRequestData {
         LookupDataResultPromisePtr promise;
         DeadlineTimerPtr timer;
+
+        void fail(Result result) {
+            cancelTimer(*timer);
+            promise->setFailed(result);
+        }
     };
 
     struct LastMessageIdRequestData {
         GetLastMessageIdResponsePromisePtr promise;
         DeadlineTimerPtr timer;
+
+        void fail(Result result) {
+            cancelTimer(*timer);
+            promise->setFailed(result);
+        }
     };
 
     struct GetSchemaRequest {
         Promise<Result, SchemaInfo> promise;
         DeadlineTimerPtr timer;
+
+        void fail(Result result) {
+            cancelTimer(*timer);
+            promise.setFailed(result);
+        }
     };
 
     /*
@@ -301,26 +320,26 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     }
 
     template <typename ConstBufferSequence, typename WriteHandler>
-    inline void asyncWrite(const ConstBufferSequence& buffers, WriteHandler handler) {
+    inline void asyncWrite(const ConstBufferSequence& buffers, WriteHandler&& handler) {
         if (isClosed()) {
             return;
         }
         if (tlsSocket_) {
-            ASIO::async_write(*tlsSocket_, buffers, ASIO::bind_executor(strand_, handler));
+            ASIO::async_write(*tlsSocket_, buffers, std::forward<WriteHandler>(handler));
         } else {
-            ASIO::async_write(*socket_, buffers, handler);
+            ASIO::async_write(*socket_, buffers, std::forward<WriteHandler>(handler));
         }
     }
 
     template <typename MutableBufferSequence, typename ReadHandler>
-    inline void asyncReceive(const MutableBufferSequence& buffers, ReadHandler handler) {
+    inline void asyncReceive(const MutableBufferSequence& buffers, ReadHandler&& handler) {
         if (isClosed()) {
             return;
         }
         if (tlsSocket_) {
-            tlsSocket_->async_read_some(buffers, ASIO::bind_executor(strand_, handler));
+            tlsSocket_->async_read_some(buffers, std::forward<ReadHandler>(handler));
         } else {
-            socket_->async_receive(buffers, handler);
+            socket_->async_receive(buffers, std::forward<ReadHandler>(handler));
         }
     }
 
@@ -341,7 +360,6 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
      */
     SocketPtr socket_;
     TlsSocketPtr tlsSocket_;
-    ASIO::strand<ASIO::io_context::executor_type> strand_;
 
     const std::string logicalAddress_;
     /*
@@ -354,7 +372,7 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     ClientConfiguration::ProxyProtocol proxyProtocol_;
 
     // Represent both endpoint of the tcp connection. eg: [client:1234 -> server:6650]
-    std::string cnxString_;
+    std::shared_ptr<std::string> cnxStringPtr_;
 
     /*
      *  indicates if async connection establishment failed
@@ -364,7 +382,8 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     SharedBuffer incomingBuffer_;
 
     Promise<Result, ClientConnectionWeakPtr> connectPromise_;
-    std::shared_ptr<PeriodicTask> connectTimeoutTask_;
+    const std::chrono::milliseconds connectTimeout_;
+    const DeadlineTimerPtr connectTimer_;
 
     typedef std::map<long, PendingRequestData> PendingRequestsMap;
     PendingRequestsMap pendingRequests_;
@@ -423,6 +442,7 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     const std::string clientVersion_;
     ConnectionPool& pool_;
     const size_t poolIndex_;
+    std::optional<std::future<void>> closeFuture_;
 
     friend class PulsarFriend;
     friend class ConsumerTest;

lib/ClientConnection.h

@@ -902,7 +902,9 @@ uint64_t Commands::serializeSingleMessagesToBatchPayload(SharedBuffer& batchPayl
         batchPayload.write(payload.data(), payload.readableBytes());
     }
 
-    return messages.back().impl_->metadata.sequence_id();
+    // Use the first message's sequence_id so that ackReceived can compute
+    // lastSequenceIdPublished_ = sequenceId + messagesCount - 1 correctly.
+    return messages.front().impl_->metadata.sequence_id();
 }
 
 Message Commands::deSerializeSingleMessageInBatch(Message& batchedMessage, int32_t batchIndex,

lib/Commands.cc

@@ -55,16 +55,43 @@ bool ConnectionPool::close() {
         return false;
     }
 
+    std::vector<ClientConnectionPtr> connectionsToClose;
+    // ClientConnection::close() will remove the connection from the pool, which is not allowed when iterating
+    // over a map, so we store the connections to close in a vector first and don't iterate the pool when
+    // closing the connections.
     std::unique_lock<std::recursive_mutex> lock(mutex_);
+    connectionsToClose.reserve(pool_.size());
+    for (auto&& kv : pool_) {
+        connectionsToClose.emplace_back(kv.second);
+    }
+    pool_.clear();
+    lock.unlock();
 
-    for (auto cnxIt = pool_.begin(); cnxIt != pool_.end(); cnxIt++) {
-        auto& cnx = cnxIt->second;
+    for (auto&& cnx : connectionsToClose) {
         if (cnx) {
-            // The 2nd argument is false because removing a value during the iteration will cause segfault
-            cnx->close(ResultDisconnected, false);
+            // Close with a fatal error to not let client retry
+            auto& future = cnx->close(ResultAlreadyClosed);
+            using namespace std::chrono_literals;
+            if (auto status = future.wait_for(5s); status != std::future_status::ready) {
+                LOG_WARN("Connection close timed out for " << cnx.get()->cnxString());
+            }
+            if (cnx.use_count() > 1) {
+                // There are some asynchronous operations that hold the reference on the connection, we should
+                // wait until them to finish. Otherwise, `io_context::stop()` will be called in
+                // `ClientImpl::shutdown()` when closing the `ExecutorServiceProvider`. Then
+                // `io_context::run()` will return and the `io_context` object will be destroyed. In this
+                // case, if there is any pending handler, it will crash.
+                for (int i = 0; i < 500 && cnx.use_count() > 1; i++) {
+                    std::this_thread::sleep_for(10ms);
+                }
+                if (cnx.use_count() > 1) {
+                    LOG_WARN("Connection still has " << (cnx.use_count() - 1)
+                                                     << " references after waiting for 5 seconds for "
+                                                     << cnx.get()->cnxString());
+                }
+            }
         }
     }
-    pool_.clear();
     return true;
 }
 
@@ -73,7 +100,7 @@ void ConnectionPool::closeAllConnectionsForNewCluster() {
     for (auto cnxIt = pool_.begin(); cnxIt != pool_.end(); cnxIt++) {
         auto& cnx = cnxIt->second;
         if (cnx) {
-            cnx->close(ResultDisconnected, false, true);
+            cnx->close(ResultDisconnected, true);
         }
     }
     pool_.clear();

lib/ConnectionPool.cc

@@ -125,8 +125,6 @@ void ExecutorService::close(long timeoutMs) {
     }
 }
 
-void ExecutorService::postWork(std::function<void(void)> task) { ASIO::post(io_context_, std::move(task)); }
-
 /////////////////////
 
 ExecutorServiceProvider::ExecutorServiceProvider(int nthreads)

lib/ExecutorService.cc

@@ -23,12 +23,16 @@
 
 #include <atomic>
 #ifdef USE_ASIO
+#include <asio/dispatch.hpp>
 #include <asio/io_context.hpp>
 #include <asio/ip/tcp.hpp>
+#include <asio/post.hpp>
 #include <asio/ssl.hpp>
 #else
+#include <boost/asio/dispatch.hpp>
 #include <boost/asio/io_context.hpp>
 #include <boost/asio/ip/tcp.hpp>
+#include <boost/asio/post.hpp>
 #include <boost/asio/ssl.hpp>
 #endif
 #include <chrono>
@@ -37,6 +41,7 @@
 #include <memory>
 #include <mutex>
 #include <thread>
+#include <utility>
 
 #include "AsioTimer.h"
 
@@ -62,7 +67,19 @@ class PULSAR_PUBLIC ExecutorService : public std::enable_shared_from_this<Execut
     TcpResolverPtr createTcpResolver();
     // throws std::runtime_error if failed
     DeadlineTimerPtr createDeadlineTimer();
-    void postWork(std::function<void(void)> task);
+
+    // Execute the task in the event loop thread asynchronously, i.e. the task will be put in the event loop
+    // queue and executed later.
+    template <typename T>
+    void postWork(T &&task) {
+        ASIO::post(io_context_, std::forward<T>(task));
+    }
+
+    // Different from `postWork`, if it's already in the event loop, execute the task immediately
+    template <typename T>
+    void dispatch(T &&task) {
+        ASIO::dispatch(io_context_, std::forward<T>(task));
+    }
 
     // See TimeoutProcessor for the semantics of the parameter.
     void close(long timeoutMs = 3000);

lib/ExecutorService.h

@@ -53,7 +53,7 @@ class PeriodicTask : public std::enable_shared_from_this<PeriodicTask> {
 
     void stop() noexcept;
 
-    void setCallback(CallbackType callback) noexcept { callback_ = callback; }
+    void setCallback(CallbackType&& callback) noexcept { callback_ = std::move(callback); }
 
     State getState() const noexcept { return state_; }
     int getPeriodMs() const noexcept { return periodMs_; }

lib/PeriodicTask.h

@@ -933,19 +933,24 @@ bool ProducerImpl::ackReceived(uint64_t sequenceId, MessageId& rawMessageId) {
         return false;
     }
 
-    uint64_t expectedSequenceId = op.sendArgs->sequenceId;
-    if (sequenceId > expectedSequenceId) {
-        LOG_WARN(getName() << "Got ack for msg " << sequenceId                        //
-                           << " expecting: " << expectedSequenceId << " queue size="  //
-                           << pendingMessagesQueue_.size() << " producer: " << producerId_);
+    const uint64_t expectedFirstSequenceId = op.sendArgs->sequenceId;
+    const uint64_t expectedLastSequenceId = expectedFirstSequenceId + op.messagesCount - 1;
+    // Broker may ack with either the first or the last sequence id of the batch.
+    if (sequenceId > expectedLastSequenceId) {
+        LOG_WARN(getName() << "Got ack for msg " << sequenceId
+                           << " expecting last: " << expectedLastSequenceId
+                           << " queue size=" << pendingMessagesQueue_.size() << " producer: " << producerId_);
         return false;
-    } else if (sequenceId < expectedSequenceId) {
+    }
+    if (sequenceId < expectedFirstSequenceId) {
         // Ignoring the ack since it's referring to a message that has already timed out.
-        LOG_DEBUG(getName() << "Got ack for timed out msg " << sequenceId  //
-                            << " -- MessageId - " << messageId << " last-seq: " << expectedSequenceId
-                            << " producer: " << producerId_);
+        LOG_DEBUG(getName() << "Got ack for timed out msg " << sequenceId << " -- MessageId - " << messageId
+                            << " first-seq: " << expectedFirstSequenceId << " producer: " << producerId_);
         return true;
     }
+    // sequenceId is in [expectedFirstSequenceId, expectedLastSequenceId]; accept as matching this op.
+    const bool brokerSentFirst = (sequenceId == expectedFirstSequenceId);
+    lastSequenceIdPublished_ = brokerSentFirst ? expectedLastSequenceId : sequenceId;
 
     // Message was persisted correctly
     LOG_DEBUG(getName() << "Received ack for msg " << sequenceId);
@@ -960,7 +965,6 @@ bool ProducerImpl::ackReceived(uint64_t sequenceId, MessageId& rawMessageId) {
     }
 
     releaseSemaphoreForSendOp(op);
-    lastSequenceIdPublished_ = sequenceId + op.messagesCount - 1;
 
     std::unique_ptr<OpSendMsg> opSendMsg{pendingMessagesQueue_.front().release()};
     pendingMessagesQueue_.pop_front();