Add tests for the race

Author: BewareMyPower

lib/ClientConnection.cc

@@ -21,6 +21,7 @@
 #include <openssl/x509.h>
 #include <pulsar/MessageIdBuilder.h>
 
+#include <chrono>
 #include <fstream>
 
 #include "AsioDefines.h"
@@ -1189,7 +1190,44 @@ Future<Result, ResponseData> ClientConnection::sendRequestWithId(const SharedBuf
     lock.unlock();
 
     LOG_DEBUG(cnxString_ << "Inserted request " << requestType << " (req_id: " << requestId << ")");
-    sendCommand(cmd);
+    if (mockingRequests_) {
+        if (auto iter = mockRequestDelays_.find(requestType); iter != mockRequestDelays_.end()) {
+            auto self = shared_from_this();
+            if (strcmp(requestType, "SEEK") == 0) {
+                // Mock the `CLOSE_CONSUMER` command sent by broker, for simplicity, disconnect all consumers
+                executor_->postWork([this, self] {
+                    std::vector<uint64_t> consumerIds;
+                    {
+                        Lock lock(mutex_);
+                        for (const auto& entry : consumers_) {
+                            if (auto consumer = entry.second.lock()) {
+                                consumerIds.push_back(consumer->getConsumerId());
+                            }
+                        }
+                    }
+                    for (auto consumerId : consumerIds) {
+                        proto::CommandCloseConsumer closeConsumerCmd;
+                        closeConsumerCmd.set_consumer_id(consumerId);
+                        self->handleCloseConsumer(closeConsumerCmd);
+                    }
+                });
+            }
+            auto timer = executor_->createDeadlineTimer();
+            timer->expires_after(std::chrono::milliseconds(iter->second));
+            LOG_INFO("Request " << requestType << " (req_id: " << requestId << ") is being delayed for "
+                                << iter->second << " ms");
+            timer->async_wait([self, cmd, requestId, timer](const ASIO_ERROR& ec) {
+                LOG_INFO("Complete request id: " << requestId);
+                proto::CommandSuccess success;
+                success.set_request_id(requestId);
+                self->handleSuccess(success);
+            });
+        } else {
+            sendCommand(cmd);
+        }
+    } else {
+        sendCommand(cmd);
+    }
     return requestData.promise.getFuture();
 }
 

lib/ClientConnection.h

@@ -123,6 +123,8 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
         Ready,
         Disconnected
     };
+    using RequestDelayType =
+        std::unordered_map<std::string /* request type */, long /* delay in milliseconds */>;
 
    public:
     typedef std::shared_ptr<ASIO::ip::tcp::socket> SocketPtr;
@@ -209,6 +211,14 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     Future<Result, SchemaInfo> newGetSchema(const std::string& topicName, const std::string& version,
                                             uint64_t requestId);
 
+    void mockRequestDelay(RequestDelayType requestDelays) {
+        if (mockingRequests_) {
+            throw new std::runtime_error("Already mocking requests");
+        }
+        mockRequestDelays_.swap(requestDelays);
+        mockingRequests_ = true;
+    }
+
    private:
     struct PendingRequestData {
         Promise<Result, ResponseData> promise;
@@ -393,6 +403,11 @@ class PULSAR_PUBLIC ClientConnection : public std::enable_shared_from_this<Clien
     DeadlineTimerPtr keepAliveTimer_;
     DeadlineTimerPtr consumerStatsRequestTimer_;
 
+    // This map should only be modified before modifying `mockingRequests_` to true for thread safety, when
+    // `mockingRequests_` is false, this field will never be accessed
+    RequestDelayType mockRequestDelays_;
+    std::atomic_bool mockingRequests_{false};
+
     void handleConsumerStatsTimeout(const ASIO_ERROR& ec, const std::vector<uint64_t>& consumerStatsRequests);
 
     void startConsumerStatsTimer(std::vector<uint64_t> consumerStatsRequests);

tests/ConsumerSeekTest.cc

@@ -24,7 +24,9 @@
 #include <string>
 
 #include "HttpHelper.h"
+#include "lib/Latch.h"
 #include "lib/LogUtils.h"
+#include "tests/PulsarFriend.h"
 
 DECLARE_LOG_OBJECT()
 
@@ -200,6 +202,23 @@ TEST_F(ConsumerSeekTest, testNoInternalConsumer) {
     ASSERT_EQ(ResultOk, consumer.seek(MessageId::earliest()));
 }
 
+// Verify the `seek` method won't be blocked forever in any order of the Subscribe response and Seek response
+TEST_F(ConsumerSeekTest, testSubscribeSeekRaces) {
+    Client client(lookupUrl);
+    Consumer consumer;
+    ASSERT_EQ(ResultOk, client.subscribe("testSubscribeSeekRaces", "sub", consumer));
+
+    for (auto&& connection : PulsarFriend::getConnections(client)) {
+        connection->mockRequestDelay({{"SUBSCRIBE", 1000}, {"SEEK", 600}});
+    }
+
+    Latch latch(1);
+    consumer.seekAsync(0L, [&latch](Result result) { latch.countdown(); });
+    ASSERT_TRUE(latch.wait(std::chrono::seconds(5)));
+
+    client.close();
+}
+
 INSTANTIATE_TEST_SUITE_P(Pulsar, ConsumerSeekTest, ::testing::Values(true, false));
 
 }  // namespace pulsar