Pipeline: use notify instead of polling for ResultQueue

dbms/src/Flash/Executor/PipelineExecutor.cpp

@@ -77,16 +77,7 @@ void PipelineExecutor::wait()
 void PipelineExecutor::consume(ResultHandler & result_handler)
 {
     assert(result_handler);
-    if (unlikely(context.isTest()))
-    {
-        // In test mode, a single query should take no more than 5 minutes to execute.
-        static std::chrono::minutes timeout(5);
-        exec_context.consumeFor(result_handler, timeout);
-    }
-    else
-    {
-        exec_context.consume(result_handler);
-    }
+    exec_context.consume(result_handler);
 }
 
 ExecutionResult PipelineExecutor::execute(ResultHandler && result_handler)

dbms/src/Flash/Executor/PipelineExecutor.h

@@ -16,7 +16,7 @@
 
 #include <Flash/Executor/PipelineExecutorContext.h>
 #include <Flash/Executor/QueryExecutor.h>
-#include <Flash/Executor/ResultQueue.h>
+#include <Flash/Executor/ResultQueue_fwd.h>
 
 namespace DB
 {

dbms/src/Flash/Executor/PipelineExecutorContext.cpp

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <Flash/Executor/PipelineExecutorContext.h>
+#include <Flash/Executor/ResultQueue.h>
 #include <Flash/Pipeline/Schedule/TaskScheduler.h>
 #include <Operators/SharedQueue.h>
 
@@ -154,6 +155,7 @@ void PipelineExecutorContext::cancel()
     if (is_cancelled.compare_exchange_strong(origin_value, true, std::memory_order_release))
     {
         cancelSharedQueues();
+        cancelResultQueueIfNeed();
         if likely (TaskScheduler::instance && !query_id.empty())
             TaskScheduler::instance->cancel(query_id, resource_group_name);
     }
@@ -163,6 +165,7 @@ ResultQueuePtr PipelineExecutorContext::toConsumeMode(size_t queue_size)
 {
     std::lock_guard lock(mu);
     RUNTIME_ASSERT(!result_queue.has_value());
+    RUNTIME_CHECK_MSG(!isCancelled(), "query has been cancelled.");
     result_queue.emplace(std::make_shared<ResultQueue>(queue_size));
     return *result_queue;
 }
@@ -185,4 +188,17 @@ void PipelineExecutorContext::cancelSharedQueues()
     for (const auto & shared_queue : tmp)
         shared_queue->cancel();
 }
+
+void PipelineExecutorContext::cancelResultQueueIfNeed()
+{
+    ResultQueue * tmp{nullptr};
+    {
+        std::lock_guard lock(mu);
+        if (isWaitMode())
+            return;
+        tmp = (*result_queue).get();
+    }
+    assert(tmp);
+    tmp->cancel();
+}
 } // namespace DB

dbms/src/Flash/Executor/PipelineExecutorContext.h

@@ -19,7 +19,7 @@
 #include <Core/AutoSpillTrigger.h>
 #include <Flash/Executor/ExecutionResult.h>
 #include <Flash/Executor/ResultHandler.h>
-#include <Flash/Executor/ResultQueue.h>
+#include <Flash/Executor/ResultQueue_fwd.h>
 #include <Flash/Pipeline/Schedule/Tasks/TaskProfileInfo.h>
 
 #include <atomic>
@@ -96,54 +96,6 @@ class PipelineExecutorContext : private boost::noncopyable
 
     void consume(ResultHandler & result_handler);
 
-    template <typename Duration>
-    void consumeFor(ResultHandler & result_handler, const Duration & timeout_duration)
-    {
-        RUNTIME_ASSERT(result_handler);
-        auto consumed_result_queue = getConsumedResultQueue();
-        bool is_timeout = false;
-        try
-        {
-            Block ret;
-            while (true)
-            {
-                auto res = consumed_result_queue->popTimeout(ret, timeout_duration);
-                if (res == MPMCQueueResult::TIMEOUT)
-                {
-                    is_timeout = true;
-                    break;
-                }
-                else if (res == MPMCQueueResult::OK)
-                {
-                    result_handler(ret);
-                }
-                else
-                {
-                    break;
-                }
-            }
-        }
-        catch (...)
-        {
-            // If result_handler throws an error, here should notify the query to terminate, and wait for the end of the query.
-            onErrorOccurred(std::current_exception());
-        }
-        if (is_timeout)
-        {
-            LOG_WARNING(log, "wait timeout");
-            onErrorOccurred(timeout_err_msg);
-            throw Exception(timeout_err_msg);
-        }
-        else
-        {
-            // In order to ensure that `decActiveRefCount` has finished calling at this point
-            // and avoid referencing the already destructed `mu` in `decActiveRefCount`.
-            std::unique_lock lock(mu);
-            cv.wait(lock, [&] { return 0 == active_ref_count; });
-        }
-        LOG_DEBUG(log, "query finished and consume done");
-    }
-
     void cancel();
 
     ALWAYS_INLINE bool isCancelled() const { return is_cancelled.load(std::memory_order_acquire); }
@@ -188,6 +140,7 @@ class PipelineExecutorContext : private boost::noncopyable
     ResultQueuePtr getConsumedResultQueue();
 
     void cancelSharedQueues();
+    void cancelResultQueueIfNeed();
 
 private:
     const String query_id;

dbms/src/Flash/Executor/ResultQueue.h

@@ -14,14 +14,35 @@
 
 #pragma once
 
-#include <Common/MPMCQueue.h>
-#include <Common/PtrHolder.h>
+#include <Common/LooseBoundedMPMCQueue.h>
 #include <Core/Block.h>
+#include <Flash/Pipeline/Schedule/Tasks/NotifyFuture.h>
 
-#include <mutex>
+#include <memory>
 
 namespace DB
 {
-using ResultQueue = MPMCQueue<Block>;
+class ResultQueue : public NotifyFuture
+{
+public:
+    explicit ResultQueue(size_t queue_size)
+        : queue(queue_size)
+    {}
+
+    // read
+    MPMCQueueResult pop(Block & block) { return queue.pop(block); }
+
+    // write
+    MPMCQueueResult push(Block && block) { return queue.push(block); }
+    MPMCQueueResult tryPush(Block && block) { return queue.tryPush(block); }
+    void registerTask(TaskPtr && task) override { queue.registerPipeWriteTask(std::move(task)); }
+
+    // finish/cancel
+    bool finish() { return queue.finish(); }
+    bool cancel() { return queue.cancel(); }
+
+private:
+    LooseBoundedMPMCQueue<Block> queue;
+};
 using ResultQueuePtr = std::shared_ptr<ResultQueue>;
 } // namespace DB

dbms/src/Flash/Executor/ResultQueue_fwd.h

@@ -0,0 +1,23 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+
+namespace DB
+{
+class ResultQueue;
+using ResultQueuePtr = std::shared_ptr<ResultQueue>;
+} // namespace DB

dbms/src/Flash/Executor/tests/gtest_pipeline_executor_context.cpp

@@ -15,6 +15,7 @@
 #include <Common/Exception.h>
 #include <Common/ThreadManager.h>
 #include <Flash/Executor/PipelineExecutorContext.h>
+#include <Flash/Executor/ResultQueue.h>
 #include <TestUtils/TiFlashTestBasic.h>
 #include <gtest/gtest.h>
 
@@ -44,29 +45,6 @@ try
 }
 CATCH
 
-TEST_F(PipelineExecutorContextTestRunner, popTimeout)
-try
-{
-    PipelineExecutorContext context;
-    context.toConsumeMode(1);
-    try
-    {
-        context.incActiveRefCount();
-        std::chrono::milliseconds timeout(10);
-        ResultHandler result_handler{[](const Block &) {
-        }};
-        context.consumeFor(result_handler, timeout);
-        GTEST_FAIL();
-    }
-    catch (DB::Exception & e)
-    {
-        GTEST_ASSERT_EQ(e.message(), PipelineExecutorContext::timeout_err_msg);
-        auto err_msg = context.getExceptionMsg();
-        ASSERT_EQ(err_msg, PipelineExecutorContext::timeout_err_msg);
-    }
-}
-CATCH
-
 TEST_F(PipelineExecutorContextTestRunner, run)
 try
 {

dbms/src/Flash/Pipeline/Pipeline.h

@@ -15,7 +15,7 @@
 #pragma once
 
 #include <Flash/Executor/ResultHandler.h>
-#include <Flash/Executor/ResultQueue.h>
+#include <Flash/Executor/ResultQueue_fwd.h>
 #include <Flash/Pipeline/Exec/PipelineExec.h>
 
 #include <deque>

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

@@ -77,15 +77,7 @@ class PipeConditionVariable
         assert(task);
         task->notify();
         assert(TaskScheduler::instance);
-        if (unlikely(task->getStatus() == ExecTaskStatus::WAITING))
-        {
-            TaskScheduler::instance->submitToWaitReactor(std::move(task));
-        }
-        else
-        {
-            assert(task->getStatus() == ExecTaskStatus::RUNNING);
-            TaskScheduler::instance->submitToCPUTaskThreadPool(std::move(task));
-        }
+        TaskScheduler::instance->submit(std::move(task));
     }
 
 private:

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

@@ -41,7 +41,7 @@ class PipelineTask
 
     ExecTaskStatus awaitImpl() override { return runAwait(); }
 
-    void notifyImpl() override { runNotify(); }
+    ExecTaskStatus notifyImpl() override { return runNotify(); }
 
     void doFinalizeImpl() override
     {

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

@@ -133,10 +133,11 @@ class PipelineTaskBase
         pipeline_exec_holder.reset();
     }
 
-    void runNotify()
+    ExecTaskStatus runNotify()
     {
         assert(pipeline_exec);
         pipeline_exec->notify();
+        return ExecTaskStatus::RUNNING;
     }
 
 private:

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

@@ -35,7 +35,7 @@ class RFWaitTask : public Task
         int max_wait_time_ms,
         RuntimeFilteList && waiting_rf_list_,
         RuntimeFilteList && ready_rf_list_)
-        : Task(exec_context_, req_id)
+        : Task(exec_context_, req_id, ExecTaskStatus::WAITING)
         , task_pool(task_pool_)
         , max_wait_time_ns(max_wait_time_ms < 0 ? 0 : 1000000UL * max_wait_time_ms)
         , waiting_rf_list(std::move(waiting_rf_list_))
@@ -75,7 +75,7 @@ class RFWaitTask : public Task
     }
 
 private:
-    ExecTaskStatus executeImpl() override { return ExecTaskStatus::WAITING; }
+    ExecTaskStatus executeImpl() override { throw Exception("unreachable"); }
 
     ExecTaskStatus awaitImpl() override
     {

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

@@ -40,7 +40,7 @@ class SimplePipelineTask
 
     ExecTaskStatus awaitImpl() override { return runAwait(); }
 
-    void notifyImpl() override { runNotify(); }
+    ExecTaskStatus notifyImpl() override { return runNotify(); }
 
     void finalizeImpl() override
     {