fix testSubscribeSeekRaces

Author: BewareMyPower

lib/ConsumerImpl.cc

@@ -238,7 +238,7 @@ Future<Result, bool> ConsumerImpl::connectionOpened(const ClientConnectionPtr& c
     cnx->registerConsumer(consumerId_, get_shared_this_ptr());
     LOG_DEBUG(cnx->cnxString() << "Registered consumer " << consumerId_);
 
-    if (duringSeek()) {
+    if (hasPendingSeek_.load(std::memory_order_acquire)) {
         ackGroupingTrackerPtr_->flushAndClean();
     }
 
@@ -269,6 +269,7 @@ Future<Result, bool> ConsumerImpl::connectionOpened(const ClientConnectionPtr& c
             } else {
                 promise.setFailed(handleResult);
             }
+            completeSeekCallback(ResultOk);
         });
 
     return promise.getFuture();
@@ -1129,19 +1130,14 @@ void ConsumerImpl::messageProcessed(Message& msg, bool track) {
  * `startMessageId_` is updated so that we can discard messages after delivery restarts.
  */
 void ConsumerImpl::clearReceiveQueue() {
-    if (duringSeek()) {
+    if (hasPendingSeek_.load(std::memory_order_acquire)) {
         if (hasSoughtByTimestamp()) {
             // Invalidate startMessageId_ so that isPriorBatchIndex and isPriorEntryIndex checks will be
             // skipped, and hasMessageAvailableAsync won't use startMessageId_ in compare.
             startMessageId_ = std::nullopt;
         } else {
             startMessageId_ = seekMessageId_.get();
         }
-        SeekStatus expected = SeekStatus::COMPLETED;
-        if (seekStatus_.compare_exchange_strong(expected, SeekStatus::NOT_STARTED)) {
-            auto seekCallback = seekCallback_.release();
-            executor_->postWork([seekCallback] { seekCallback(ResultOk); });
-        }
         return;
     } else if (subscriptionMode_ == Commands::SubscriptionModeDurable) {
         return;
@@ -1554,7 +1550,9 @@ void ConsumerImpl::seekAsync(const MessageId& msgId, const ResultCallback& callb
     }
 
     const auto requestId = newRequestId();
-    seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, msgId), SeekArg{msgId}, callback);
+    auto nonNullCallback = (callback != nullptr) ? callback : [](Result) {};
+    seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, msgId), SeekArg{msgId},
+                      std::move(nonNullCallback));
 }
 
 void ConsumerImpl::seekAsync(uint64_t timestamp, const ResultCallback& callback) {
@@ -1568,8 +1566,9 @@ void ConsumerImpl::seekAsync(uint64_t timestamp, const ResultCallback& callback)
     }
 
     const auto requestId = newRequestId();
+    auto nonNullCallback = (callback != nullptr) ? callback : [](Result) {};
     seekAsyncInternal(requestId, Commands::newSeek(consumerId_, requestId, timestamp), SeekArg{timestamp},
-                      callback);
+                      std::move(nonNullCallback));
 }
 
 bool ConsumerImpl::isReadCompacted() { return readCompacted_; }
@@ -1727,18 +1726,17 @@ bool ConsumerImpl::isConnected() const { return !getCnx().expired() && state_ ==
 uint64_t ConsumerImpl::getNumberOfConnectedConsumer() { return isConnected() ? 1 : 0; }
 
 void ConsumerImpl::seekAsyncInternal(long requestId, const SharedBuffer& seek, const SeekArg& seekArg,
-                                     const ResultCallback& callback) {
+                                     ResultCallback&& callback) {
     ClientConnectionPtr cnx = getCnx().lock();
     if (!cnx) {
         LOG_ERROR(getName() << " Client Connection not ready for Consumer");
         callback(ResultNotConnected);
         return;
     }
 
-    auto expected = SeekStatus::NOT_STARTED;
-    if (!seekStatus_.compare_exchange_strong(expected, SeekStatus::IN_PROGRESS)) {
-        LOG_ERROR(getName() << " attempted to seek " << seekArg << " when the status is "
-                            << static_cast<int>(expected));
+    auto expected = false;
+    if (hasPendingSeek_.compare_exchange_strong(expected, true)) {
+        LOG_ERROR(getName() << " attempted to seek " << seekArg << " when there is a pending seek");
         callback(ResultNotAllowedError);
         return;
     }
@@ -1750,8 +1748,7 @@ void ConsumerImpl::seekAsyncInternal(long requestId, const SharedBuffer& seek, c
         seekMessageId_ = *boost::get<MessageId>(&seekArg);
         hasSoughtByTimestamp_.store(false, std::memory_order_release);
     }
-    seekStatus_ = SeekStatus::IN_PROGRESS;
-    seekCallback_ = callback;
+    seekCallback_ = std::move(callback);
     LOG_INFO(getName() << " Seeking subscription to " << seekArg);
 
     auto weakSelf = weak_from_this();
@@ -1771,20 +1768,17 @@ void ConsumerImpl::seekAsyncInternal(long requestId, const SharedBuffer& seek, c
                 Lock lock(mutexForMessageId_);
                 lastDequedMessageId_ = MessageId::earliest();
                 lock.unlock();
-                if (getCnx().expired()) {
-                    // It's during reconnection, complete the seek future after connection is established
-                    seekStatus_ = SeekStatus::COMPLETED;
-                } else {
+
+                if (!getCnx().expired()) {
                     if (!hasSoughtByTimestamp()) {
                         startMessageId_ = seekMessageId_.get();
                     }
-                    seekCallback_.release()(result);
-                }
+                    completeSeekCallback(result);
+                }  // else: complete the seek future after connection is established
             } else {
                 LOG_ERROR(getName() << "Failed to seek: " << result);
                 seekMessageId_ = originalSeekMessageId;
-                seekStatus_ = SeekStatus::NOT_STARTED;
-                seekCallback_.release()(result);
+                completeSeekCallback(result);
             }
         });
 }

lib/ConsumerImpl.h

@@ -21,6 +21,7 @@
 
 #include <pulsar/Reader.h>
 
+#include <atomic>
 #include <boost/variant.hpp>
 #include <cstdint>
 #include <functional>
@@ -77,13 +78,6 @@ const static std::string SYSTEM_PROPERTY_REAL_TOPIC = "REAL_TOPIC";
 const static std::string PROPERTY_ORIGIN_MESSAGE_ID = "ORIGIN_MESSAGE_ID";
 const static std::string DLQ_GROUP_TOPIC_SUFFIX = "-DLQ";
 
-enum class SeekStatus : std::uint8_t
-{
-    NOT_STARTED,
-    IN_PROGRESS,
-    COMPLETED
-};
-
 class ConsumerImpl : public ConsumerImplBase {
    public:
     ConsumerImpl(const ClientImplPtr& client, const std::string& topic, const std::string& subscriptionName,
@@ -230,7 +224,13 @@ class ConsumerImpl : public ConsumerImplBase {
     }
 
     void seekAsyncInternal(long requestId, const SharedBuffer& seek, const SeekArg& seekArg,
-                           const ResultCallback& callback);
+                           ResultCallback&& callback);
+    void completeSeekCallback(Result result) {
+        if (auto callback = seekCallback_.release()) {
+            callback(result);
+        }
+        hasPendingSeek_.store(false, std::memory_order_release);
+    }
     void processPossibleToDLQ(const MessageId& messageId, const ProcessDLQCallBack& cb);
 
     std::mutex mutexForReceiveWithZeroQueueSize;
@@ -274,14 +274,13 @@ class ConsumerImpl : public ConsumerImplBase {
     MessageId lastDequedMessageId_{MessageId::earliest()};
     MessageId lastMessageIdInBroker_{MessageId::earliest()};
 
-    std::atomic<SeekStatus> seekStatus_{SeekStatus::NOT_STARTED};
     Synchronized<ResultCallback> seekCallback_{[](Result) {}};
     Synchronized<optional<MessageId>> startMessageId_;
+    std::atomic_bool hasPendingSeek_{false};
     Synchronized<MessageId> seekMessageId_{MessageId::earliest()};
     std::atomic<bool> hasSoughtByTimestamp_{false};
 
     bool hasSoughtByTimestamp() const { return hasSoughtByTimestamp_.load(std::memory_order_acquire); }
-    bool duringSeek() const { return seekStatus_ != SeekStatus::NOT_STARTED; }
 
     class ChunkedMessageCtx {
        public:

lib/Synchronized.h

@@ -41,6 +41,12 @@ class Synchronized {
         return *this;
     }
 
+    Synchronized& operator=(T&& value) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        value_ = value;
+        return *this;
+    }
+
    private:
     T value_;
     mutable std::mutex mutex_;