Pipeline: use notify instead of polling for ExchangeReceiver

dbms/src/Flash/Coprocessor/DAGContext.h

@@ -351,6 +351,8 @@ class DAGContext
     UInt64 getConnectionID() const { return connection_id; }
     const String & getConnectionAlias() const { return connection_alias; }
 
+    MPPReceiverSetPtr getMPPReceiverSet() const { return mpp_receiver_set; }
+
 public:
     DAGRequest dag_request;
     /// Some existing code inherited from Clickhouse assume that each query must have a valid query string and query ast,
@@ -443,6 +445,7 @@ class DAGContext
     /// warning_count is the actual warning count during the entire execution
     std::atomic<UInt64> warning_count;
 
+    // `mpp_receiver_set` is always set by `MPPTask` and is used later.
     MPPReceiverSetPtr mpp_receiver_set;
     std::vector<CoprocessorReaderPtr> coprocessor_readers;
     /// vector of SubqueriesForSets(such as join build subquery).

dbms/src/Flash/Executor/PipelineExecutorContext.cpp

@@ -15,6 +15,7 @@
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Executor/PipelineExecutorContext.h>
 #include <Flash/Executor/ResultQueue.h>
+#include <Flash/Mpp/MPPReceiverSet.h>
 #include <Flash/Mpp/MPPTunnelSet.h>
 #include <Flash/Mpp/Utils.h>
 #include <Flash/Pipeline/Schedule/TaskScheduler.h>
@@ -178,9 +179,12 @@ void PipelineExecutorContext::cancel()
         cancelSharedQueues();
         if (likely(dag_context))
         {
-            // Cancel the tunnel_set here to prevent pipeline tasks waiting in the WAIT_FOR_NOTIFY state from never being notified.
+            // Cancel the tunnel_set and mpp_receiver_set here to prevent
+            // pipeline tasks waiting in the WAIT_FOR_NOTIFY state from never being notified.
             if (dag_context->tunnel_set)
                 dag_context->tunnel_set->close(getTrimmedErrMsg(), false);
+            if (auto mpp_receiver_set = dag_context->getMPPReceiverSet(); mpp_receiver_set)
+                mpp_receiver_set->cancel();
         }
         cancelResultQueueIfNeed();
         if likely (TaskScheduler::instance && !query_id.empty())

dbms/src/Flash/Mpp/LocalRequestHandler.h

@@ -42,7 +42,6 @@ struct LocalRequestHandler
 
     bool isWritable() const { return msg_queue->isWritable(); }
 
-    void registerPipeReadTask(TaskPtr && task) const { msg_queue->registerPipeReadTask(std::move(task)); }
     void registerPipeWriteTask(TaskPtr && task) const { msg_queue->registerPipeWriteTask(std::move(task)); }
 
     void writeDone(bool meet_error, const String & local_err_msg) const

dbms/src/Flash/Mpp/MPPTunnel.cpp

@@ -403,7 +403,7 @@ WaitResult MPPTunnel::waitForWritable() const
         RUNTIME_CHECK_MSG(tunnel_sender != nullptr, "write to tunnel {} which is already closed.", tunnel_id);
         if (!tunnel_sender->isWritable())
         {
-            setNotifyFuture(tunnel_sender);
+            setNotifyFuture(tunnel_sender.get());
             return WaitResult::WaitForNotify;
         }
         return WaitResult::Ready;

dbms/src/Flash/Mpp/ReceivedMessageQueue.cpp

@@ -104,7 +104,7 @@ ReceivedMessageQueue::ReceivedMessageQueue(
               : std::function<void(const ReceivedMessagePtr &)>([this](const ReceivedMessagePtr & element) {
                     for (size_t i = 0; i < fine_grained_channel_size; ++i)
                     {
-                        auto result = msg_channels_for_fine_grained_shuffle[i]->forcePush(element);
+                        auto result = msg_channels_for_fine_grained_shuffle[i].forcePush(element);
                         RUNTIME_CHECK_MSG(result == MPMCQueueResult::OK, "push to fine grained channel must success");
                     }
                 }))
@@ -114,9 +114,7 @@ ReceivedMessageQueue::ReceivedMessageQueue(
         assert(fine_grained_channel_size > 0);
         msg_channels_for_fine_grained_shuffle.reserve(fine_grained_channel_size);
         for (size_t i = 0; i < fine_grained_channel_size; ++i)
-            /// these are unbounded queues
-            msg_channels_for_fine_grained_shuffle.push_back(
-                std::make_shared<LooseBoundedMPMCQueue<ReceivedMessagePtr>>(std::numeric_limits<size_t>::max()));
+            msg_channels_for_fine_grained_shuffle.emplace_back();
     }
 }
 
@@ -127,9 +125,9 @@ MPMCQueueResult ReceivedMessageQueue::pop(size_t stream_id, ReceivedMessagePtr &
     if (fine_grained_channel_size > 0)
     {
         if constexpr (need_wait)
-            res = msg_channels_for_fine_grained_shuffle[stream_id]->pop(recv_msg);
+            res = msg_channels_for_fine_grained_shuffle[stream_id].pop(recv_msg);
         else
-            res = msg_channels_for_fine_grained_shuffle[stream_id]->tryPop(recv_msg);
+            res = msg_channels_for_fine_grained_shuffle[stream_id].tryPop(recv_msg);
 
         if (res == MPMCQueueResult::OK)
         {
@@ -151,6 +149,15 @@ MPMCQueueResult ReceivedMessageQueue::pop(size_t stream_id, ReceivedMessagePtr &
 #else
                 grpc_recv_queue.tryDequeue();
 #endif
+                ExchangeReceiverMetric::subDataSizeMetric(*data_size_in_queue, recv_msg->getPacket().ByteSizeLong());
+            }
+        }
+        else
+        {
+            if constexpr (!need_wait)
+            {
+                if (res == MPMCQueueResult::EMPTY)
+                    setNotifyFuture(&msg_channels_for_fine_grained_shuffle[stream_id]);
             }
         }
     }
@@ -160,13 +167,20 @@ MPMCQueueResult ReceivedMessageQueue::pop(size_t stream_id, ReceivedMessagePtr &
             res = grpc_recv_queue.pop(recv_msg);
         else
             res = grpc_recv_queue.tryPop(recv_msg);
-    }
 
-    if (res == MPMCQueueResult::OK)
-    {
-        ExchangeReceiverMetric::subDataSizeMetric(*data_size_in_queue, recv_msg->getPacket().ByteSizeLong());
+        if (res == MPMCQueueResult::OK)
+        {
+            ExchangeReceiverMetric::subDataSizeMetric(*data_size_in_queue, recv_msg->getPacket().ByteSizeLong());
+        }
+        else
+        {
+            if constexpr (!need_wait)
+            {
+                if (res == MPMCQueueResult::EMPTY)
+                    setNotifyFuture(&grpc_recv_queue);
+            }
+        }
     }
-
     return res;
 }
 

dbms/src/Flash/Mpp/ReceivedMessageQueue.h

@@ -21,8 +21,10 @@
 #include <Common/TiFlashMetrics.h>
 #include <Flash/Mpp/ReceivedMessage.h>
 #include <Flash/Mpp/TrackedMppDataPacket.h>
+#include <Flash/Pipeline/Schedule/Tasks/NotifyFuture.h>
 
 #include <memory>
+#include <utility>
 
 namespace DB
 {
@@ -55,6 +57,65 @@ enum class ReceiverMode
     Async
 };
 
+class GRPCNotifyQueue : public NotifyFuture
+{
+public:
+    template <typename... Args>
+    explicit GRPCNotifyQueue(const LoggerPtr & log_, Args &&... args)
+        : queue(log_, std::forward<Args>(args)...)
+    {}
+
+    void registerTask(TaskPtr && task) override { queue.registerPipeReadTask(std::move(task)); }
+
+    MPMCQueueResult pop(ReceivedMessagePtr & data) { return queue.pop(data); }
+
+    MPMCQueueResult tryPop(ReceivedMessagePtr & data) { return queue.tryPop(data); }
+
+    MPMCQueueResult forcePush(ReceivedMessagePtr && data) { return queue.forcePush(std::move(data)); }
+
+    MPMCQueueResult push(ReceivedMessagePtr && data) { return queue.push(std::move(data)); }
+
+    MPMCQueueResult tryDequeue() { return queue.tryDequeue(); }
+
+    MPMCQueueResult pushWithTag(ReceivedMessagePtr && data, GRPCKickTag * new_tag)
+    {
+        return queue.pushWithTag(std::move(data), new_tag);
+    }
+
+    void setKickFuncForTest(GRPCKickFunc && func) { queue.setKickFuncForTest(std::move(func)); }
+
+    bool finish() { return queue.finish(); }
+    bool cancel() { return queue.cancel(); }
+
+    bool isWritable() const { return queue.isWritable(); }
+
+    void registerPipeWriteTask(TaskPtr && task) { queue.registerPipeWriteTask(std::move(task)); }
+
+private:
+    GRPCRecvQueue<ReceivedMessagePtr> queue;
+};
+
+class MSGChannel : public NotifyFuture
+{
+public:
+    void registerTask(TaskPtr && task) override { queue_ref.registerPipeReadTask(std::move(task)); }
+
+    MPMCQueueResult pop(ReceivedMessagePtr & data) { return queue_ref.pop(data); }
+
+    MPMCQueueResult tryPop(ReceivedMessagePtr & data) { return queue_ref.tryPop(data); }
+
+    MPMCQueueResult forcePush(const ReceivedMessagePtr & data) { return queue_ref.forcePush(data); }
+
+    bool finish() { return queue_ref.finish(); }
+    bool cancel() { return queue_ref.cancel(); }
+
+private:
+    using QueueImpl = LooseBoundedMPMCQueue<ReceivedMessagePtr>;
+    // these are unbounded queues.
+    std::unique_ptr<QueueImpl> queue = std::make_unique<QueueImpl>(std::numeric_limits<size_t>::max());
+    QueueImpl & queue_ref = *queue;
+};
+
 class ReceivedMessageQueue
 {
 public:
@@ -86,20 +147,19 @@ class ReceivedMessageQueue
         grpc_recv_queue.finish();
         /// msg_channels_for_fine_grained_shuffle must be finished after msg_channel is finished
         for (auto & channel : msg_channels_for_fine_grained_shuffle)
-            channel->finish();
+            channel.finish();
     }
 
     void cancel()
     {
         grpc_recv_queue.cancel();
         /// msg_channels_for_fine_grained_shuffle must be cancelled after msg_channel is cancelled
         for (auto & channel : msg_channels_for_fine_grained_shuffle)
-            channel->cancel();
+            channel.cancel();
     }
 
     bool isWritable() const { return grpc_recv_queue.isWritable(); }
 
-    void registerPipeReadTask(TaskPtr && task) { grpc_recv_queue.registerPipeReadTask(std::move(task)); }
     void registerPipeWriteTask(TaskPtr && task) { grpc_recv_queue.registerPipeWriteTask(std::move(task)); }
 
 #ifndef DBMS_PUBLIC_GTEST
@@ -118,8 +178,8 @@ class ReceivedMessageQueue
     /// write: the writer first write the msg to msg_channel/grpc_recv_queue, if write success, then write msg to msg_channels_for_fine_grained_shuffle
     /// read: the reader read msg from msg_channels_for_fine_grained_shuffle, and reduce the `remaining_consumers` in msg, if `remaining_consumers` is 0, then
     ///       remove the msg from msg_channel/grpc_recv_queue
-    std::vector<std::shared_ptr<LooseBoundedMPMCQueue<ReceivedMessagePtr>>> msg_channels_for_fine_grained_shuffle;
-    GRPCRecvQueue<ReceivedMessagePtr> grpc_recv_queue;
+    std::vector<MSGChannel> msg_channels_for_fine_grained_shuffle;
+    GRPCNotifyQueue grpc_recv_queue;
 };
 
 } // namespace DB

dbms/src/Flash/Pipeline/Schedule/Tasks/Impls/StreamRestoreTask.cpp

@@ -48,7 +48,7 @@ ExecTaskStatus StreamRestoreTask::tryFlush()
         t_block.clear();
         return ExecTaskStatus::IO_IN;
     case MPMCQueueResult::FULL:
-        setNotifyFuture(sink);
+        setNotifyFuture(sink.get());
         return ExecTaskStatus::WAIT_FOR_NOTIFY;
     case MPMCQueueResult::CANCELLED:
         return ExecTaskStatus::CANCELLED;

dbms/src/Flash/Pipeline/Schedule/Tasks/NotifyFuture.cpp

@@ -16,23 +16,23 @@
 
 namespace DB
 {
-thread_local NotifyFuturePtr current_notify_future = nullptr;
+thread_local NotifyFuture * current_notify_future = nullptr;
 
-void setNotifyFuture(NotifyFuturePtr new_future)
+void setNotifyFuture(NotifyFuture * new_future)
 {
     assert(current_notify_future == nullptr);
     current_notify_future = std::move(new_future);
 }
 
 void clearNotifyFuture()
 {
-    current_notify_future.reset();
+    current_notify_future = nullptr;
 }
 
 void registerTaskToFuture(TaskPtr && task)
 {
     assert(current_notify_future != nullptr);
     current_notify_future->registerTask(std::move(task));
-    current_notify_future.reset();
+    current_notify_future = nullptr;
 }
 } // namespace DB

dbms/src/Flash/Pipeline/Schedule/Tasks/NotifyFuture.h

@@ -24,11 +24,10 @@ struct NotifyFuture
     virtual ~NotifyFuture() = default;
     virtual void registerTask(TaskPtr && task) = 0;
 };
-using NotifyFuturePtr = std::shared_ptr<NotifyFuture>;
 
-extern thread_local NotifyFuturePtr current_notify_future;
+extern thread_local NotifyFuture * current_notify_future;
 
-void setNotifyFuture(NotifyFuturePtr new_future);
+void setNotifyFuture(NotifyFuture * new_future);
 void clearNotifyFuture();
 void registerTaskToFuture(TaskPtr && task);
 

dbms/src/Flash/Pipeline/Schedule/Tasks/Task.h

@@ -26,12 +26,12 @@ namespace DB
  *           CANCELLED/ERROR/FINISHED
  *                      ▲
  *                      │
- *         ┌───────────────────────────────┐
- *         │     ┌──►RUNNING◄──┐           │
- * INIT───►│     │             │           │
- *         │     ▼             ▼           │
- *         │ WATITING◄────────►IO_IN/OUT   │
- *         └───────────────────────────────┘
+ * ┌───────────────────────────────────────────────┐
+ * │     ┌──────────►RUNNING◄──────────┐           │
+ * │     │                             │           │
+ * │     ▼                             ▼           │
+ * │ WATITING/WAIT_FOR_NOTIFY◄────────►IO_IN/OUT   │
+ * └───────────────────────────────────────────────┘
  */
 enum class ExecTaskStatus
 {