Reduce unnecessary getLastMessageId RPC in hasMessageAvailable

Author: BewareMyPower

lib/ConsumerImpl.cc

@@ -961,6 +961,9 @@ Result ConsumerImpl::fetchSingleMessageFromBroker(Message& msg) {
                 // Can't use break here else it may trigger a race with connection opened.
 
                 localLock.unlock();
+                Lock lock(mutexForMessageId_);
+                lastDequedMessageId_ = msg.getMessageId();
+                lock.unlock();
                 msg = interceptors_->beforeConsume(Consumer(shared_from_this()), msg);
                 return ResultOk;
             }
@@ -1107,9 +1110,7 @@ void ConsumerImpl::messageProcessed(Message& msg, bool track) {
  */
 void ConsumerImpl::clearReceiveQueue() {
     if (duringSeek()) {
-        if (!hasSoughtByTimestamp_.load(std::memory_order_acquire)) {
-            startMessageId_ = seekMessageId_.get();
-        }
+        trySetStartMessageIdToSeekMessageId();
         SeekStatus expected = SeekStatus::COMPLETED;
         if (seekStatus_.compare_exchange_strong(expected, SeekStatus::NOT_STARTED)) {
             auto seekCallback = seekCallback_.release();
@@ -1561,14 +1562,18 @@ void ConsumerImpl::seekAsync(uint64_t timestamp, ResultCallback callback) {
 bool ConsumerImpl::isReadCompacted() { return readCompacted_; }
 
 void ConsumerImpl::hasMessageAvailableAsync(HasMessageAvailableCallback callback) {
+    if (!incomingMessages_.empty()) {
+        callback(ResultOk, true);
+        return;
+    }
     bool compareMarkDeletePosition;
     {
         std::lock_guard<std::mutex> lock{mutexForMessageId_};
         compareMarkDeletePosition =
             (lastDequedMessageId_ == MessageId::earliest()) &&
             (startMessageId_.get().value_or(MessageId::earliest()) == MessageId::latest());
     }
-    if (compareMarkDeletePosition || hasSoughtByTimestamp_.load(std::memory_order_acquire)) {
+    if (compareMarkDeletePosition || hasSoughtByTimestamp_.load()) {
         auto self = get_shared_this_ptr();
         getLastMessageIdAsync([self, callback](Result result, const GetLastMessageIdResponse& response) {
             if (result != ResultOk) {
@@ -1587,8 +1592,7 @@ void ConsumerImpl::hasMessageAvailableAsync(HasMessageAvailableCallback callback
                     callback(ResultOk, false);
                 }
             };
-            if (self->config_.isStartMessageIdInclusive() &&
-                !self->hasSoughtByTimestamp_.load(std::memory_order_acquire)) {
+            if (self->config_.isStartMessageIdInclusive() && !self->hasSoughtByTimestamp_.load()) {
                 self->seekAsync(response.getLastMessageId(), [callback, handleResponse](Result result) {
                     if (result != ResultOk) {
                         callback(result, {});
@@ -1723,7 +1727,7 @@ void ConsumerImpl::seekAsyncInternal(long requestId, SharedBuffer seek, const Se
 
     const auto originalSeekMessageId = seekMessageId_.get();
     if (boost::get<uint64_t>(&seekArg)) {
-        hasSoughtByTimestamp_.store(true, std::memory_order_release);
+        hasSoughtByTimestamp_.store(true);
     } else {
         seekMessageId_ = *boost::get<MessageId>(&seekArg);
     }
@@ -1752,16 +1756,15 @@ void ConsumerImpl::seekAsyncInternal(long requestId, SharedBuffer seek, const Se
                     // It's during reconnection, complete the seek future after connection is established
                     seekStatus_ = SeekStatus::COMPLETED;
                 } else {
-                    if (!hasSoughtByTimestamp_.load(std::memory_order_acquire)) {
-                        startMessageId_ = seekMessageId_.get();
-                    }
+                    trySetStartMessageIdToSeekMessageId();
                     seekCallback_.release()(result);
                 }
             } else {
                 LOG_ERROR(getName() << "Failed to seek: " << result);
                 seekMessageId_ = originalSeekMessageId;
                 seekStatus_ = SeekStatus::NOT_STARTED;
                 seekCallback_.release()(result);
+                hasSoughtByTimestamp_.store(false);
             }
         });
 }

lib/ConsumerImpl.h

@@ -372,6 +372,14 @@ class ConsumerImpl : public ConsumerImplBase {
         }
     }
 
+    void trySetStartMessageIdToSeekMessageId() {
+        bool expected = true;
+        if (!hasSoughtByTimestamp_.compare_exchange_strong(expected, false)) {
+            startMessageId_ = seekMessageId_.get();
+        }  // else: we have sought by timestamp, we should not set startMessageId_ to seekMessageId_
+           // because it is not a valid message id.
+    }
+
     friend class PulsarFriend;
     friend class MultiTopicsConsumerImpl;
 

tests/ReaderTest.cc

@@ -21,7 +21,10 @@
 #include <pulsar/Reader.h>
 #include <time.h>
 
+#include <chrono>
+#include <future>
 #include <string>
+#include <thread>
 
 #include "HttpHelper.h"
 #include "PulsarFriend.h"
@@ -917,5 +920,50 @@ TEST_F(ReaderSeekTest, testSeekInclusiveChunkMessage) {
     assertStartMessageId(false, secondMsgId);
 }
 
+// https://github.com/apache/pulsar-client-cpp/pull/422 introduces a regression that after seeking by
+// timestamp, each hasMessageAvailable() call will require a getLastMessageId RPC. This test is added to
+// verify the regression is fixed.
+TEST_F(ReaderSeekTest, testHasMessageAvailableCache) {
+    const auto topic = "test-has-message-available-cache-" + std::to_string(time(nullptr));
+
+    Reader reader;
+    auto hasMessageAvailableCallbackThreadId = [&reader] {
+        std::promise<std::thread::id> threadIdPromise;
+        auto callback = [&threadIdPromise](Result, bool) {
+            threadIdPromise.set_value(std::this_thread::get_id());
+        };
+        reader.hasMessageAvailableAsync(callback);
+        return threadIdPromise.get_future().get();
+    };
+
+    Producer producer;
+    ASSERT_EQ(ResultOk, client.createProducer(topic, producer));
+    producer.send(MessageBuilder().setContent("msg-0").build());
+    producer.send(MessageBuilder().setContent("msg-1").build());
+
+    // Case 1: when `incomingMessages_` is empty, `hasMessageAvailableAsync` will complete immediately if
+    // `lastDequedMessageId_` and `lastMessageIdInBroker_` have been updated and there is a message available.
+    ReaderConfiguration conf;
+    conf.setReceiverQueueSize(0);
+    ASSERT_EQ(ResultOk, client.createReader(topic, MessageId::latest(), conf, reader));
+    ASSERT_EQ(ResultOk, reader.seek(0L));
+    Message msg;
+    ASSERT_EQ(ResultOk, reader.readNext(msg));
+    // Now, `lastDequedMessageId_` is updated to the id of `msg`
+    bool available;
+    ASSERT_EQ(ResultOk, reader.hasMessageAvailable(available));
+    // Now the `lastMessageIdInBroker_` is updated
+    ASSERT_EQ(hasMessageAvailableCallbackThreadId(), std::this_thread::get_id());
+
+    // Case 2: when `incomingMessages_` is not empty, `hasMessageAvailableAsync` will complete immediately
+    ASSERT_EQ(ResultOk, client.createReader(topic, MessageId::latest(), {}, reader));
+    ASSERT_EQ(ResultOk, reader.seek(0L));
+    ASSERT_EQ(ResultOk, reader.readNext(msg));
+    waitUntil(std::chrono::seconds(3),
+              [&reader] { return PulsarFriend::getConsumer(reader)->getNumOfPrefetchedMessages() > 0; });
+    // Now, `incomingMessages_` is not empty
+    ASSERT_EQ(hasMessageAvailableCallbackThreadId(), std::this_thread::get_id());
+}
+
 INSTANTIATE_TEST_SUITE_P(Pulsar, ReaderTest, ::testing::Values(true, false));
 INSTANTIATE_TEST_SUITE_P(Pulsar, ReaderSeekTest, ::testing::Values(true, false));