Add non-blocking functions for MPMCQueue (#5311)

close #5310

dbms/src/Common/MPMCQueue.h

@@ -74,56 +74,80 @@ class MPMCQueue
             destruct(getObj(read_pos));
     }
 
-    /// Block util:
+    /// Block until:
     /// 1. Pop succeeds with a valid T: return true.
     /// 2. The queue is cancelled or finished: return false.
-    bool pop(T & obj)
+    ALWAYS_INLINE bool pop(T & obj)
     {
-        return popObj(obj);
+        return popObj<true>(obj);
     }
 
-    /// Besides all conditions mentioned at `pop`, `tryPop` will return false if `timeout` is exceeded.
+    /// Besides all conditions mentioned at `pop`, `popTimeout` will return false if `timeout` is exceeded.
     template <typename Duration>
-    bool tryPop(T & obj, const Duration & timeout)
+    ALWAYS_INLINE bool popTimeout(T & obj, const Duration & timeout)
     {
         /// std::condition_variable::wait_until will always use system_clock.
         auto deadline = std::chrono::system_clock::now() + timeout;
-        return popObj(obj, &deadline);
+        return popObj<true>(obj, &deadline);
     }
 
-    /// Block util:
+    /// Non-blocking function.
+    /// Return true if pop succeed.
+    /// else return false.
+    ALWAYS_INLINE bool tryPop(T & obj)
+    {
+        return popObj<false>(obj);
+    }
+
+    /// Block until:
     /// 1. Push succeeds and return true.
     /// 2. The queue is cancelled and return false.
     /// 3. The queue has finished and return false.
     template <typename U>
     ALWAYS_INLINE bool push(U && u)
     {
-        return pushObj(std::forward<U>(u));
+        return pushObj<true>(std::forward<U>(u));
     }
 
-    /// Besides all conditions mentioned at `push`, `tryPush` will return false if `timeout` is exceeded.
+    /// Besides all conditions mentioned at `push`, `pushTimeout` will return false if `timeout` is exceeded.
     template <typename U, typename Duration>
-    ALWAYS_INLINE bool tryPush(U && u, const Duration & timeout)
+    ALWAYS_INLINE bool pushTimeout(U && u, const Duration & timeout)
     {
         /// std::condition_variable::wait_until will always use system_clock.
         auto deadline = std::chrono::system_clock::now() + timeout;
-        return pushObj(std::forward<U>(u), &deadline);
+        return pushObj<true>(std::forward<U>(u), &deadline);
+    }
+
+    /// Non-blocking function.
+    /// Return true if push succeed.
+    /// else return false.
+    template <typename U>
+    ALWAYS_INLINE bool tryPush(U && u)
+    {
+        return pushObj<false>(std::forward<U>(u));
     }
 
     /// The same as `push` except it will construct the object in place.
     template <typename... Args>
     ALWAYS_INLINE bool emplace(Args &&... args)
     {
-        return emplaceObj(nullptr, std::forward<Args>(args)...);
+        return emplaceObj<true>(nullptr, std::forward<Args>(args)...);
     }
 
-    /// The same as `tryPush` except it will construct the object in place.
+    /// The same as `pushTimeout` except it will construct the object in place.
     template <typename... Args, typename Duration>
-    ALWAYS_INLINE bool tryEmplace(Args &&... args, const Duration & timeout)
+    ALWAYS_INLINE bool emplaceTimeout(Args &&... args, const Duration & timeout)
     {
         /// std::condition_variable::wait_until will always use system_clock.
         auto deadline = std::chrono::system_clock::now() + timeout;
-        return emplaceObj(&deadline, std::forward<Args>(args)...);
+        return emplaceObj<true>(&deadline, std::forward<Args>(args)...);
+    }
+
+    /// The same as `tryPush` except it will construct the object in place.
+    template <typename... Args>
+    ALWAYS_INLINE bool tryEmplace(Args &&... args)
+    {
+        return emplaceObj<false>(nullptr, std::forward<Args>(args)...);
     }
 
     /// Cancel a NORMAL queue will wake up all blocking readers and writers.
@@ -233,22 +257,25 @@ class MPMCQueue
         }
     }
 
-    bool popObj(T & res, const TimePoint * deadline = nullptr)
+    template <bool need_wait>
+    bool popObj(T & res, [[maybe_unused]] const TimePoint * deadline = nullptr)
     {
 #ifdef __APPLE__
         WaitingNode node;
 #else
         thread_local WaitingNode node;
 #endif
         {
-            /// read_pos < write_pos means the queue isn't empty
-            auto pred = [&] {
-                return read_pos < write_pos || !isNormal();
-            };
-
             std::unique_lock lock(mu);
 
-            wait(lock, reader_head, node, pred, deadline);
+            if constexpr (need_wait)
+            {
+                /// read_pos < write_pos means the queue isn't empty
+                auto pred = [&] {
+                    return read_pos < write_pos || !isNormal();
+                };
+                wait(lock, reader_head, node, pred, deadline);
+            }
 
             if (!isCancelled() && read_pos < write_pos)
             {
@@ -272,21 +299,23 @@ class MPMCQueue
         return false;
     }
 
-    template <typename F>
-    bool assignObj(const TimePoint * deadline, F && assigner)
+    template <bool need_wait, typename F>
+    bool assignObj([[maybe_unused]] const TimePoint * deadline, F && assigner)
     {
 #ifdef __APPLE__
         WaitingNode node;
 #else
         thread_local WaitingNode node;
 #endif
-        auto pred = [&] {
-            return write_pos - read_pos < capacity || !isNormal();
-        };
-
         std::unique_lock lock(mu);
 
-        wait(lock, writer_head, node, pred, deadline);
+        if constexpr (need_wait)
+        {
+            auto pred = [&] {
+                return write_pos - read_pos < capacity || !isNormal();
+            };
+            wait(lock, writer_head, node, pred, deadline);
+        }
 
         /// double check status after potential wait
         /// check write_pos because timeouted will also reach here.
@@ -305,16 +334,16 @@ class MPMCQueue
         return false;
     }
 
-    template <typename U>
+    template <bool need_wait, typename U>
     ALWAYS_INLINE bool pushObj(U && u, const TimePoint * deadline = nullptr)
     {
-        return assignObj(deadline, [&](void * addr) { new (addr) T(std::forward<U>(u)); });
+        return assignObj<need_wait>(deadline, [&](void * addr) { new (addr) T(std::forward<U>(u)); });
     }
 
-    template <typename... Args>
+    template <bool need_wait, typename... Args>
     ALWAYS_INLINE bool emplaceObj(const TimePoint * deadline, Args &&... args)
     {
-        return assignObj(deadline, [&](void * addr) { new (addr) T(std::forward<Args>(args)...); });
+        return assignObj<need_wait>(deadline, [&](void * addr) { new (addr) T(std::forward<Args>(args)...); });
     }
 
     ALWAYS_INLINE bool isNormal() const

dbms/src/Common/tests/gtest_mpmc_queue.cpp

@@ -98,12 +98,14 @@ class MPMCQueueTest : public ::testing::Test
     void testCannotTryPush(MPMCQueue<T> & queue)
     {
         auto old_size = queue.size();
-        auto res = queue.tryPush(ValueHelper<T>::make(-1), std::chrono::microseconds(1));
-        auto new_size = queue.size();
-        if (res)
+        bool ok1 = queue.tryPush(ValueHelper<T>::make(-1));
+        auto new_size1 = queue.size();
+        bool ok2 = queue.pushTimeout(ValueHelper<T>::make(-1), std::chrono::microseconds(1));
+        auto new_size2 = queue.size();
+        if (ok1 || ok2)
             throw TiFlashTestException("Should push fail");
-        if (old_size != new_size)
-            throw TiFlashTestException(fmt::format("Size changed from {} to {} without push", old_size, new_size));
+        if (old_size != new_size1 || old_size != new_size2)
+            throw TiFlashTestException(fmt::format("Size changed from {} to {} and {} without push", old_size, new_size1, new_size2));
     }
 
     template <typename T>
@@ -124,12 +126,14 @@ class MPMCQueueTest : public ::testing::Test
     {
         auto old_size = queue.size();
         T res;
-        bool ok = queue.tryPop(res, std::chrono::microseconds(1));
-        auto new_size = queue.size();
-        if (ok)
+        bool ok1 = queue.tryPop(res);
+        auto new_size1 = queue.size();
+        bool ok2 = queue.popTimeout(res, std::chrono::microseconds(1));
+        auto new_size2 = queue.size();
+        if (ok1 || ok2)
             throw TiFlashTestException("Should pop fail");
-        if (old_size != new_size)
-            throw TiFlashTestException(fmt::format("Size changed from {} to {} without pop", old_size, new_size));
+        if (old_size != new_size1 || old_size != new_size2)
+            throw TiFlashTestException(fmt::format("Size changed from {} to {} and {} without pop", old_size, new_size1, new_size2));
     }
 
     template <typename T>
@@ -474,7 +478,6 @@ class MPMCQueueTest : public ::testing::Test
             throwOrMove(std::move(rhs));
         }
 
-
         ThrowInjectable & operator=(ThrowInjectable && rhs)
         {
             if (this != &rhs)

dbms/src/Common/tests/mpmc_queue_perftest.cpp

@@ -87,7 +87,7 @@ struct Helper<MPMCQueue<T>>
     template <typename U>
     static void pushOneTo(MPMCQueue<T> & queue, U && data)
     {
-        queue.tryPush(std::forward<U>(data), std::chrono::milliseconds(1));
+        queue.pushTimeout(std::forward<U>(data), std::chrono::milliseconds(1));
     }
 };
 

dbms/src/Flash/Mpp/ExchangeReceiver.cpp

@@ -424,7 +424,7 @@ void ExchangeReceiverBase<RPCContext>::reactor(const std::vector<Request> & asyn
         for (Int32 i = 0; i < check_waiting_requests_freq; ++i)
         {
             AsyncHandler * handler = nullptr;
-            if (unlikely(!ready_requests.tryPop(handler, timeout)))
+            if (unlikely(!ready_requests.popTimeout(handler, timeout)))
                 break;
 
             handler->handle();