Pipeline: use notify instead of polling for SharedQueue

dbms/src/Common/LooseBoundedMPMCQueue.h

@@ -15,7 +15,9 @@
 #pragma once
 
 #include <Common/MPMCQueue.h>
+#include <Flash/Pipeline/Schedule/Tasks/PipeConditionVariable.h>
 
+#include <condition_variable>
 #include <deque>
 
 namespace DB
@@ -48,46 +50,81 @@ class LooseBoundedMPMCQueue
         , push_callback(std::move(push_callback_))
     {}
 
-    /// blocking function.
-    /// Just like MPMCQueue::push.
-    template <typename U>
-    MPMCQueueResult push(U && data)
+    void registerPipeReadTask(TaskPtr && task)
     {
-        std::unique_lock lock(mu);
-        writer_head.wait(lock, [&] { return !isFullWithoutLock() || (unlikely(status != MPMCQueueStatus::NORMAL)); });
+        {
+            std::lock_guard lock(mu);
+            if (queue.empty() && status == MPMCQueueStatus::NORMAL)
+            {
+                pipe_reader_cv.registerTask(std::move(task));
+                return;
+            }
+        }
+        PipeConditionVariable::notifyTaskDirectly(std::move(task));
+    }
 
-        if ((likely(status == MPMCQueueStatus::NORMAL)) && !isFullWithoutLock())
+    void registerPipeWriteTask(TaskPtr && task)
+    {
         {
-            pushFront(std::forward<U>(data));
-            return MPMCQueueResult::OK;
+            std::lock_guard lock(mu);
+            if (isFullWithoutLock() && status == MPMCQueueStatus::NORMAL)
+            {
+                pipe_writer_cv.registerTask(std::move(task));
+                return;
+            }
         }
+        PipeConditionVariable::notifyTaskDirectly(std::move(task));
+    }
 
-        switch (status)
+    /// blocking function.
+    /// Just like MPMCQueue::push.
+    template <typename U>
+    MPMCQueueResult push(U && data)
+    {
+        bool notify_writer{false};
         {
-        case MPMCQueueStatus::NORMAL:
-            return MPMCQueueResult::FULL;
-        case MPMCQueueStatus::CANCELLED:
-            return MPMCQueueResult::CANCELLED;
-        case MPMCQueueStatus::FINISHED:
-            return MPMCQueueResult::FINISHED;
+            std::unique_lock lock(mu);
+            writer_cv.wait(lock, [&] { return !isFullWithoutLock() || (unlikely(status != MPMCQueueStatus::NORMAL)); });
+
+            if unlikely (status == MPMCQueueStatus::CANCELLED)
+                return MPMCQueueResult::CANCELLED;
+            else if unlikely (status == MPMCQueueStatus::FINISHED)
+                return MPMCQueueResult::FINISHED;
+            else
+            {
+                assert(status == MPMCQueueStatus::NORMAL);
+                assert(!isFullWithoutLock());
+                notify_writer = pushFront(std::forward<U>(data));
+            }
         }
+        if (notify_writer)
+            notifyOneWriter();
+        notifyOneReader();
+        return MPMCQueueResult::OK;
     }
 
     /// Just like MPMCQueue::tryPush.
     template <typename U>
     MPMCQueueResult tryPush(U && data)
     {
-        std::lock_guard lock(mu);
-
-        if unlikely (status == MPMCQueueStatus::CANCELLED)
-            return MPMCQueueResult::CANCELLED;
-        if unlikely (status == MPMCQueueStatus::FINISHED)
-            return MPMCQueueResult::FINISHED;
-
-        if (isFullWithoutLock())
-            return MPMCQueueResult::FULL;
-
-        pushFront(std::forward<U>(data));
+        bool notify_writer{false};
+        {
+            std::lock_guard lock(mu);
+
+            if unlikely (status == MPMCQueueStatus::CANCELLED)
+                return MPMCQueueResult::CANCELLED;
+            else if unlikely (status == MPMCQueueStatus::FINISHED)
+                return MPMCQueueResult::FINISHED;
+            else
+            {
+                if (isFullWithoutLock())
+                    return MPMCQueueResult::FULL;
+                notify_writer = pushFront(std::forward<U>(data));
+            }
+        }
+        if (notify_writer)
+            notifyOneWriter();
+        notifyOneReader();
         return MPMCQueueResult::OK;
     }
 
@@ -97,64 +134,74 @@ class LooseBoundedMPMCQueue
     template <typename U>
     MPMCQueueResult forcePush(U && data)
     {
-        std::lock_guard lock(mu);
-
-        if unlikely (status == MPMCQueueStatus::CANCELLED)
-            return MPMCQueueResult::CANCELLED;
-        if unlikely (status == MPMCQueueStatus::FINISHED)
-            return MPMCQueueResult::FINISHED;
-
-        pushFront(std::forward<U>(data));
+        bool notify_writer{false};
+        {
+            std::lock_guard lock(mu);
+
+            if unlikely (status == MPMCQueueStatus::CANCELLED)
+                return MPMCQueueResult::CANCELLED;
+            else if unlikely (status == MPMCQueueStatus::FINISHED)
+                return MPMCQueueResult::FINISHED;
+            else
+                notify_writer = pushFront(std::forward<U>(data));
+        }
+        if (notify_writer)
+            notifyOneWriter();
+        notifyOneReader();
         return MPMCQueueResult::OK;
     }
 
     MPMCQueueResult pop(T & data)
     {
-        std::unique_lock lock(mu);
-        reader_head.wait(lock, [&] { return !queue.empty() || (unlikely(status != MPMCQueueStatus::NORMAL)); });
-
-        if ((likely(status != MPMCQueueStatus::CANCELLED)) && !queue.empty())
         {
-            data = popBack();
-            return MPMCQueueResult::OK;
-        }
+            std::unique_lock lock(mu);
+            reader_cv.wait(lock, [&] { return !queue.empty() || (unlikely(status != MPMCQueueStatus::NORMAL)); });
 
-        switch (status)
-        {
-        case MPMCQueueStatus::NORMAL:
-            return MPMCQueueResult::EMPTY;
-        case MPMCQueueStatus::CANCELLED:
-            return MPMCQueueResult::CANCELLED;
-        case MPMCQueueStatus::FINISHED:
-            return MPMCQueueResult::FINISHED;
+            if unlikely (status == MPMCQueueStatus::CANCELLED)
+                return MPMCQueueResult::CANCELLED;
+
+            if (queue.empty())
+            {
+                assert(status == MPMCQueueStatus::FINISHED);
+                return MPMCQueueResult::FINISHED;
+            }
+
+            data = popBack();
         }
+        notifyOneWriter();
+        return MPMCQueueResult::OK;
     }
 
     MPMCQueueResult tryPop(T & data)
     {
-        std::lock_guard lock(mu);
-
-        if unlikely (status == MPMCQueueStatus::CANCELLED)
-            return MPMCQueueResult::CANCELLED;
+        {
+            std::lock_guard lock(mu);
+            if unlikely (status == MPMCQueueStatus::CANCELLED)
+                return MPMCQueueResult::CANCELLED;
 
-        if (queue.empty())
-            return status == MPMCQueueStatus::NORMAL ? MPMCQueueResult::EMPTY : MPMCQueueResult::FINISHED;
+            if (queue.empty())
+                return status == MPMCQueueStatus::NORMAL ? MPMCQueueResult::EMPTY : MPMCQueueResult::FINISHED;
 
-        data = popBack();
+            data = popBack();
+        }
+        notifyOneWriter();
         return MPMCQueueResult::OK;
     }
 
     MPMCQueueResult tryDequeue()
     {
-        std::lock_guard lock(mu);
+        {
+            std::lock_guard lock(mu);
 
-        if unlikely (status == MPMCQueueStatus::CANCELLED)
-            return MPMCQueueResult::CANCELLED;
+            if unlikely (status == MPMCQueueStatus::CANCELLED)
+                return MPMCQueueResult::CANCELLED;
 
-        if (queue.empty())
-            return status == MPMCQueueStatus::NORMAL ? MPMCQueueResult::EMPTY : MPMCQueueResult::FINISHED;
+            if (queue.empty())
+                return status == MPMCQueueStatus::NORMAL ? MPMCQueueResult::EMPTY : MPMCQueueResult::FINISHED;
 
-        popBack();
+            popBack();
+        }
+        notifyOneWriter();
         return MPMCQueueResult::OK;
     }
 
@@ -208,6 +255,27 @@ class LooseBoundedMPMCQueue
     }
 
 private:
+    void notifyOneReader()
+    {
+        reader_cv.notify_one();
+        pipe_reader_cv.notifyOne();
+    }
+
+    void notifyOneWriter()
+    {
+        writer_cv.notify_one();
+        pipe_writer_cv.notifyOne();
+    }
+
+    void notifyAll()
+    {
+        reader_cv.notify_all();
+        pipe_reader_cv.notifyAll();
+
+        writer_cv.notify_all();
+        pipe_writer_cv.notifyAll();
+    }
+
     bool isFullWithoutLock() const
     {
         assert(current_auxiliary_memory_usage >= 0);
@@ -218,15 +286,18 @@ class LooseBoundedMPMCQueue
     template <typename FF>
     ALWAYS_INLINE bool changeStatus(FF && ff)
     {
-        std::lock_guard lock(mu);
-        if likely (status == MPMCQueueStatus::NORMAL)
+        bool ret{false};
         {
-            ff();
-            reader_head.notifyAll();
-            writer_head.notifyAll();
-            return true;
+            std::lock_guard lock(mu);
+            if likely (status == MPMCQueueStatus::NORMAL)
+            {
+                ff();
+                ret = true;
+            }
         }
-        return false;
+        if (ret)
+            notifyAll();
+        return ret;
     }
 
     ALWAYS_INLINE T popBack()
@@ -235,12 +306,12 @@ class LooseBoundedMPMCQueue
         queue.pop_back();
         current_auxiliary_memory_usage -= element.memory_usage;
         assert(!queue.empty() || current_auxiliary_memory_usage == 0);
-        writer_head.notifyNext();
         return element.data;
     }
 
+    // If returns true, then notify writers afterward; if false, no need to notify writers.
     template <typename U>
-    ALWAYS_INLINE void pushFront(U && data)
+    ALWAYS_INLINE bool pushFront(U && data)
     {
         Int64 memory_usage = get_auxiliary_memory_usage(data);
         queue.emplace_front(std::forward<U>(data), memory_usage);
@@ -249,7 +320,7 @@ class LooseBoundedMPMCQueue
         {
             push_callback(queue.front().data);
         }
-        reader_head.notifyNext();
+
         /// consider a case that the queue capacity is 2, the max_auxiliary_memory_usage is 100,
         /// T1: a writer write an object with size 100
         /// T2: two writers(w2, w3) try to write, but all blocked because of the max_auxiliary_memory_usage
@@ -260,8 +331,9 @@ class LooseBoundedMPMCQueue
         /// 1. there is another reader
         /// 2. there is another writer
         /// if we notify the writer if the queue is not full here, w3 can write immediately
-        if (capacity_limits.max_bytes != std::numeric_limits<Int64>::max() && !isFullWithoutLock())
-            writer_head.notifyNext();
+        ///
+        /// return true means that writer should be notified.
+        return capacity_limits.max_bytes != std::numeric_limits<Int64>::max() && !isFullWithoutLock();
     }
 
 private:
@@ -286,8 +358,11 @@ class LooseBoundedMPMCQueue
     const PushCallback push_callback;
     Int64 current_auxiliary_memory_usage = 0;
 
-    MPMCQueueDetail::WaitingNode reader_head;
-    MPMCQueueDetail::WaitingNode writer_head;
+    std::condition_variable reader_cv;
+    std::condition_variable writer_cv;
+
+    PipeConditionVariable pipe_reader_cv;
+    PipeConditionVariable pipe_writer_cv;
 
     MPMCQueueStatus status = MPMCQueueStatus::NORMAL;
     String cancel_reason;

dbms/src/Flash/Coprocessor/InterpreterUtils.cpp

@@ -99,16 +99,24 @@ void restoreConcurrency(
 {
     if (concurrency > 1 && group_builder.concurrency() == 1)
     {
-        auto shared_queue = SharedQueue::build(1, concurrency, max_buffered_bytes);
+        // Doesn't use `auto [shared_queue_sink_holder, shared_queue_source_holder]` just to make c++ compiler happy.
+        SharedQueueSinkHolderPtr shared_queue_sink_holder;
+        SharedQueueSourceHolderPtr shared_queue_source_holder;
+        std::tie(shared_queue_sink_holder, shared_queue_source_holder)
+            = SharedQueue::build(exec_context, 1, concurrency, max_buffered_bytes);
         // sink op of builder must be empty.
         group_builder.transform([&](auto & builder) {
-            builder.setSinkOp(std::make_unique<SharedQueueSinkOp>(exec_context, log->identifier(), shared_queue));
+            builder.setSinkOp(
+                std::make_unique<SharedQueueSinkOp>(exec_context, log->identifier(), shared_queue_sink_holder));
         });
         auto cur_header = group_builder.getCurrentHeader();
         group_builder.addGroup();
         for (size_t i = 0; i < concurrency; ++i)
-            group_builder.addConcurrency(
-                std::make_unique<SharedQueueSourceOp>(exec_context, log->identifier(), cur_header, shared_queue));
+            group_builder.addConcurrency(std::make_unique<SharedQueueSourceOp>(
+                exec_context,
+                log->identifier(),
+                cur_header,
+                shared_queue_source_holder));
     }
 }
 
@@ -120,14 +128,22 @@ void executeUnion(
 {
     if (group_builder.concurrency() > 1)
     {
-        auto shared_queue = SharedQueue::build(group_builder.concurrency(), 1, max_buffered_bytes);
+        // Doesn't use `auto [shared_queue_sink_holder, shared_queue_source_holder]` just to make c++ compiler happy.
+        SharedQueueSinkHolderPtr shared_queue_sink_holder;
+        SharedQueueSourceHolderPtr shared_queue_source_holder;
+        std::tie(shared_queue_sink_holder, shared_queue_source_holder)
+            = SharedQueue::build(exec_context, group_builder.concurrency(), 1, max_buffered_bytes);
         group_builder.transform([&](auto & builder) {
-            builder.setSinkOp(std::make_unique<SharedQueueSinkOp>(exec_context, log->identifier(), shared_queue));
+            builder.setSinkOp(
+                std::make_unique<SharedQueueSinkOp>(exec_context, log->identifier(), shared_queue_sink_holder));
         });
         auto cur_header = group_builder.getCurrentHeader();
         group_builder.addGroup();
-        group_builder.addConcurrency(
-            std::make_unique<SharedQueueSourceOp>(exec_context, log->identifier(), cur_header, shared_queue));
+        group_builder.addConcurrency(std::make_unique<SharedQueueSourceOp>(
+            exec_context,
+            log->identifier(),
+            cur_header,
+            shared_queue_source_holder));
     }
 }