Support cancel on ActiveSegmentReadTaskQueue

Signed-off-by: JaySon-Huang <tshent@qq.com>

Author: JaySon-Huang

dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp

@@ -1380,13 +1380,15 @@ void DAGStorageInterpreter::buildLocalExec(
         assert(!table_query_infos.empty()); // check at start of this function
         auto read_queue = table_query_infos.begin()->second.read_queue;
         read_queue->finishQueueIfEmpty();
+        exec_context.addStorageTaskQueue(read_queue);
     }
     else
     {
         for (const auto & table_query_info : table_query_infos)
         {
             const SelectQueryInfo & query_info = table_query_info.second;
             query_info.read_queue->finishQueueIfEmpty();
+            exec_context.addStorageTaskQueue(query_info.read_queue);
         }
     }
 

dbms/src/Flash/Executor/PipelineExecutorContext.cpp

@@ -22,6 +22,7 @@
 #include <Flash/Pipeline/Schedule/Tasks/OneTimeNotifyFuture.h>
 #include <Operators/CTE.h>
 #include <Operators/SharedQueue.h>
+#include <Storages/DeltaMerge/ReadThread/ActiveSegmentReadTaskQueue.h>
 
 #include <exception>
 
@@ -178,7 +179,7 @@ void PipelineExecutorContext::cancel()
     bool origin_value = false;
     if (is_cancelled.compare_exchange_strong(origin_value, true, std::memory_order_release))
     {
-        // TODO: Add ActiveSegmentReadTaskQueue here to cancel storage layer read tasks
+        cancelStorageTaskQueues();
         cancelSharedQueues();
         cancelOneTimeFutures();
         if (likely(dag_context))
@@ -235,6 +236,25 @@ void PipelineExecutorContext::cancelSharedQueues()
         shared_queue->cancel();
 }
 
+void PipelineExecutorContext::addStorageTaskQueue(const DM::ActiveSegmentReadTaskQueuePtr & storage_task_queue)
+{
+    std::lock_guard lock(mu);
+    RUNTIME_CHECK_MSG(!isCancelled(), "query has been cancelled.");
+    assert(storage_task_queue);
+    storage_task_queues.push_back(storage_task_queue);
+}
+
+void PipelineExecutorContext::cancelStorageTaskQueues()
+{
+    std::vector<DM::ActiveSegmentReadTaskQueuePtr> tmp;
+    {
+        std::lock_guard lock(mu);
+        std::swap(tmp, storage_task_queues);
+    }
+    for (const auto & storage_task_queue : tmp)
+        storage_task_queue->finishQueue();
+}
+
 void PipelineExecutorContext::addOneTimeFuture(const OneTimeNotifyFuturePtr & future)
 {
     std::lock_guard lock(mu);

dbms/src/Flash/Executor/PipelineExecutorContext.h

@@ -40,6 +40,13 @@ class OneTimeNotifyFuture;
 using OneTimeNotifyFuturePtr = std::shared_ptr<OneTimeNotifyFuture>;
 class DAGContext;
 
+namespace DM
+{
+class ActiveSegmentReadTaskQueue;
+using ActiveSegmentReadTaskQueuePtr = std::shared_ptr<ActiveSegmentReadTaskQueue>;
+} // namespace DM
+
+
 class PipelineExecutorContext : private boost::noncopyable
 {
 public:
@@ -146,6 +153,8 @@ class PipelineExecutorContext : private boost::noncopyable
 
     void addOneTimeFuture(const OneTimeNotifyFuturePtr & future);
 
+    void addStorageTaskQueue(const DM::ActiveSegmentReadTaskQueuePtr & storage_task_queue);
+
 private:
     bool setExceptionPtr(const std::exception_ptr & exception_ptr_);
 
@@ -160,6 +169,8 @@ class PipelineExecutorContext : private boost::noncopyable
 
     void cancelOneTimeFutures();
 
+    void cancelStorageTaskQueues();
+
     void cancelResultQueueIfNeed();
 
 private:
@@ -196,5 +207,7 @@ class PipelineExecutorContext : private boost::noncopyable
     std::vector<SharedQueuePtr> shared_queues;
 
     std::vector<OneTimeNotifyFuturePtr> one_time_futures;
+
+    std::vector<DM::ActiveSegmentReadTaskQueuePtr> storage_task_queues;
 };
 } // namespace DB

dbms/src/Storages/StorageDisaggregatedRemote.cpp

@@ -620,7 +620,9 @@ std::variant<DM::Remote::RNWorkersPtr, DM::SegmentReadTaskPoolPtr> StorageDisagg
 
     if (enable_read_thread)
     {
-        TableID table_id = table_scan.getLogicalTableID(); // FIXME:
+        // Now for StorageDisaggregated, we create only one SegmentReadTaskPool for segments from all partitions.
+        // Use the logical table id as a workaround.
+        TableID table_id = table_scan.getLogicalTableID(); 
         return std::make_shared<DM::SegmentReadTaskPool>(
             read_queue,
             extra_table_id_index,
@@ -690,7 +692,8 @@ void StorageDisaggregated::buildRemoteSegmentInputStreams(
     size_t num_streams,
     DAGPipeline & pipeline)
 {
-    // Share the read queue among all inputstreams // TODO: share for partition tables under disagg
+    // Share the read queue among all inputstreams. Note that for StorageDisaggregated,
+    // now we create only one SegmentReadTaskPool for segment from all partitions.
     auto read_queue = std::make_shared<DM::ActiveSegmentReadTaskQueue>(num_streams, log);
     // Build the input streams to read blocks from remote segments
     auto [column_defines, extra_table_id_index] = genColumnDefinesForDisaggregatedRead(table_scan);
@@ -763,8 +766,10 @@ void StorageDisaggregated::buildRemoteSegmentSourceOps(
     DM::SegmentReadTasks && read_tasks,
     size_t num_streams)
 {
-    // Share the read queue among all source ops // TODO: share for partition tables under disagg
+    // Share the read queue among all source ops. Note that for StorageDisaggregated,
+    // now we create only one SegmentReadTaskPool for segment from all partitions.
     auto read_queue = std::make_shared<DM::ActiveSegmentReadTaskQueue>(num_streams, log);
+    exec_context.addStorageTaskQueue(read_queue);
     // Build the input streams to read blocks from remote segments
     auto [column_defines, extra_table_id_index] = genColumnDefinesForDisaggregatedRead(table_scan);
     auto packed_read_tasks = packSegmentReadTasks(