track incomplete requests

Author: BewareMyPower

lib/MockServer.h

@@ -19,6 +19,7 @@
 #pragma once
 
 #include <initializer_list>
+#include <memory>
 #include <mutex>
 #include <unordered_map>
 #include <vector>
@@ -31,7 +32,7 @@
 
 namespace pulsar {
 
-class MockServer {
+class MockServer : public std::enable_shared_from_this<MockServer> {
    public:
     using RequestDelayType = std::unordered_map<std::string, long /* delay in milliseconds */>;
 
@@ -72,18 +73,25 @@ class MockServer {
             }
             long delayMs = iter->second;
             auto timer = connection->executor_->createDeadlineTimer();
+            auto key = request + std::to_string(requestId);
+            pendingTimers_[key] = timer;
             timer->expires_from_now(std::chrono::milliseconds(delayMs));
-            timer->async_wait([connection, requestId, request, timer](const auto& ec) {
+
+            auto self = shared_from_this();
+            timer->async_wait([this, self, key, connection, requestId, timer](const auto& ec) {
+                {
+                    std::lock_guard<std::mutex> lock(mutex_);
+                    pendingTimers_.erase(key);
+                }
                 if (ec) {
-                    LOG_INFO("Timer cancelled for request " << request << " with id " << requestId);
+                    LOG_INFO("Timer cancelled for request " << key);
                     return;
                 }
                 if (connection->isClosed()) {
-                    LOG_INFO("Connection is closed, not completing request " << request << " with id "
-                                                                             << requestId);
+                    LOG_INFO("Connection is closed, not completing request " << key);
                     return;
                 }
-                LOG_INFO("Completing delayed request " << request << " with id " << requestId);
+                LOG_INFO("Completing delayed request " << key);
                 proto::CommandSuccess success;
                 success.set_request_id(requestId);
                 connection->handleSuccess(success);
@@ -94,9 +102,25 @@ class MockServer {
         }
     }
 
+    // Return the number of pending timers cancelled
+    auto close() {
+        std::lock_guard<std::mutex> lock(mutex_);
+        auto result = pendingTimers_.size();
+        for (auto&& kv : pendingTimers_) {
+            try {
+                LOG_INFO("Cancelling timer for " << kv.first);
+                kv.second->cancel();
+            } catch (...) {
+                LOG_WARN("Failed to cancel timer for " << kv.first);
+            }
+        }
+        return result;
+    }
+
    private:
     mutable std::mutex mutex_;
     std::unordered_map<std::string, long> requestDelays_;
+    std::unordered_map<std::string, DeadlineTimerPtr> pendingTimers_;
     ClientConnectionWeakPtr connection_;
 
     DECLARE_LOG_OBJECT()

tests/ConsumerSeekTest.cc

@@ -209,7 +209,12 @@ TEST_F(ConsumerSeekTest, testNoInternalConsumer) {
 static void assertSeekWithTimeout(Consumer& consumer) {
     using namespace std::chrono_literals;
     auto promise = std::make_shared<std::promise<Result>>();
-    consumer.seekAsync(0L, [promise](Result result) { promise->set_value(result); });
+    std::weak_ptr<std::promise<Result>> weakPromise = promise;
+    consumer.seekAsync(0L, [weakPromise](Result result) {
+        if (auto promise = weakPromise.lock()) {
+            promise->set_value(result);
+        }
+    });
     auto future = promise->get_future();
     ASSERT_EQ(future.wait_for(5s), std::future_status::ready);
     ASSERT_EQ(future.get(), ResultOk);
@@ -231,13 +236,14 @@ TEST_F(ConsumerSeekTest, testSubscribeSeekRaces) {
     mockServer->setRequestDelay({{"SUBSCRIBE", 500}, {"SEEK", 1000}});
     assertSeekWithTimeout(consumer);
 
+    ASSERT_EQ(mockServer->close(), 0);
     client.close();
 }
 
 TEST_F(ConsumerSeekTest, testReconnectionSlow) {
     Client client(lookupUrl, ClientConfiguration().setInitialBackoffIntervalMs(500));
     Consumer consumer;
-    ASSERT_EQ(ResultOk, client.subscribe("testSubscribeSeekRaces", "sub", consumer));
+    ASSERT_EQ(ResultOk, client.subscribe("testReconnectionSlow", "sub", consumer));
 
     auto connection = *PulsarFriend::getConnections(client).begin();
     auto mockServer = std::make_shared<MockServer>(connection);
@@ -247,6 +253,7 @@ TEST_F(ConsumerSeekTest, testReconnectionSlow) {
     mockServer->setRequestDelay({{"SEEK", 100}});
     assertSeekWithTimeout(consumer);
 
+    ASSERT_EQ(mockServer->close(), 0);
     client.close();
 }