join be aware of cancel signal (#9450) (#9452)

close #9430

Co-authored-by: xufei <xufei@pingcap.com>
Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

dbms/src/DataStreams/AggregatingBlockInputStream.cpp

@@ -32,7 +32,7 @@ Block AggregatingBlockInputStream::readImpl()
         executed = true;
         AggregatedDataVariantsPtr data_variants = std::make_shared<AggregatedDataVariants>();
 
-        Aggregator::CancellationHook hook = [&]() {
+        CancellationHook hook = [&]() {
             return this->isCancelled();
         };
         aggregator.setCancellationHook(hook);

dbms/src/DataStreams/HashJoinProbeExec.h

@@ -36,8 +36,6 @@ class HashJoinProbeExec : public std::enable_shared_from_this<HashJoinProbeExec>
         const BlockInputStreamPtr & probe_stream,
         size_t max_block_size);
 
-    using CancellationHook = std::function<bool()>;
-
     HashJoinProbeExec(
         const String & req_id,
         const JoinPtr & join_,

dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp

@@ -71,7 +71,7 @@ Block ParallelAggregatingBlockInputStream::readImpl()
 {
     if (!executed)
     {
-        Aggregator::CancellationHook hook = [&]() {
+        CancellationHook hook = [&]() {
             return this->isCancelled();
         };
         aggregator.setCancellationHook(hook);

dbms/src/Flash/Mpp/MPPTask.cpp

@@ -517,6 +517,9 @@ void MPPTask::runImpl()
         GET_METRIC(tiflash_coprocessor_request_duration_seconds, type_run_mpp_task).Observe(stopwatch.elapsedSeconds());
     });
 
+    // set cancellation hook
+    context->setCancellationHook([this] { return is_cancelled.load(); });
+
     String err_msg;
     try
     {
@@ -746,7 +749,7 @@ void MPPTask::abort(const String & message, AbortType abort_type)
         if (previous_status == FINISHED || previous_status == CANCELLED || previous_status == FAILED)
         {
             LOG_WARNING(log, "task already in {} state", magic_enum::enum_name(previous_status));
-            return;
+            break;
         }
         else if (previous_status == INITIALIZING && switchStatus(INITIALIZING, next_task_status))
         {
@@ -755,7 +758,7 @@ void MPPTask::abort(const String & message, AbortType abort_type)
             /// so just close all tunnels here
             abortTunnels("", false);
             LOG_WARNING(log, "Finish abort task from uninitialized");
-            return;
+            break;
         }
         else if (previous_status == RUNNING && switchStatus(RUNNING, next_task_status))
         {
@@ -769,9 +772,10 @@ void MPPTask::abort(const String & message, AbortType abort_type)
             scheduleThisTask(ScheduleState::FAILED);
             /// runImpl is running, leave remaining work to runImpl
             LOG_WARNING(log, "Finish abort task from running");
-            return;
+            break;
         }
     }
+    is_cancelled = true;
 }
 
 bool MPPTask::switchStatus(TaskStatus from, TaskStatus to)

dbms/src/Flash/Mpp/MPPTask.h

@@ -177,6 +177,7 @@ class MPPTask
 
     MPPTaskManager * manager;
     std::atomic<bool> is_registered{false};
+    std::atomic<bool> is_cancelled{false};
 
     MPPTaskScheduleEntry schedule_entry;
 

dbms/src/Flash/Planner/Plans/PhysicalJoin.cpp

@@ -238,6 +238,7 @@ void PhysicalJoin::probeSideTransform(DAGPipeline & probe_pipeline, Context & co
             settings.max_block_size);
         stream->setExtraInfo(join_probe_extra_info);
     }
+    join_ptr->setCancellationHook([&] { return context.isCancelled(); });
 }
 
 void PhysicalJoin::buildSideTransform(DAGPipeline & build_pipeline, Context & context, size_t max_streams)

dbms/src/Flash/Planner/Plans/PhysicalJoinProbe.cpp

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <Flash/Coprocessor/InterpreterUtils.h>
+#include <Flash/Executor/PipelineExecutorContext.h>
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
 #include <Flash/Planner/Plans/PhysicalJoinProbe.h>
 #include <Interpreters/Context.h>
@@ -56,6 +57,7 @@ void PhysicalJoinProbe::buildPipelineExecGroupImpl(
             max_block_size,
             input_header));
     });
+    join_ptr->setCancellationHook([&]() { return exec_context.isCancelled(); });
     join_ptr.reset();
 }
 } // namespace DB

dbms/src/Interpreters/Aggregator.h

@@ -32,6 +32,7 @@
 #include <Interpreters/AggSpillContext.h>
 #include <Interpreters/AggregateDescription.h>
 #include <Interpreters/AggregationCommon.h>
+#include <Interpreters/CancellationHook.h>
 #include <TiDB/Collation/Collator.h>
 #include <common/StringRef.h>
 #include <common/logger_useful.h>
@@ -1199,8 +1200,6 @@ class Aggregator
       */
     Blocks convertBlockToTwoLevel(const Block & block);
 
-    using CancellationHook = std::function<bool()>;
-
     /** Set a function that checks whether the current task can be aborted.
       */
     void setCancellationHook(CancellationHook cancellation_hook);

dbms/src/Interpreters/CancellationHook.h

@@ -0,0 +1,22 @@
+// 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 <functional>
+
+namespace DB
+{
+using CancellationHook = std::function<bool()>;
+} // namespace DB

dbms/src/Interpreters/Context.h

@@ -19,6 +19,7 @@
 #include <Core/Types.h>
 #include <Debug/MockServerInfo.h>
 #include <IO/FileProvider/FileProvider_fwd.h>
+#include <Interpreters/CancellationHook.h>
 #include <Interpreters/ClientInfo.h>
 #include <Interpreters/Context_fwd.h>
 #include <Interpreters/Settings.h>
@@ -178,6 +179,9 @@ class Context
     TimezoneInfo timezone_info;
 
     DAGContext * dag_context = nullptr;
+    CancellationHook is_cancelled{[]() {
+        return false;
+    }};
     using DatabasePtr = std::shared_ptr<IDatabase>;
     using Databases = std::map<String, std::shared_ptr<IDatabase>>;
     /// Use copy constructor or createGlobal() instead
@@ -235,8 +239,8 @@ class Context
     /// Compute and set actual user settings, client_info.current_user should be set
     void calculateUserSettings();
 
-    ClientInfo & getClientInfo() { return client_info; };
-    const ClientInfo & getClientInfo() const { return client_info; };
+    ClientInfo & getClientInfo() { return client_info; }
+    const ClientInfo & getClientInfo() const { return client_info; }
 
     void setQuota(
         const String & name,
@@ -373,6 +377,9 @@ class Context
     void setDAGContext(DAGContext * dag_context);
     DAGContext * getDAGContext() const;
 
+    bool isCancelled() const { return is_cancelled(); }
+    void setCancellationHook(CancellationHook cancellation_hook) { is_cancelled = cancellation_hook; }
+
     /// List all queries.
     ProcessList & getProcessList();
     const ProcessList & getProcessList() const;
@@ -499,8 +506,8 @@ class Context
 
     SharedQueriesPtr getSharedQueries();
 
-    const TimezoneInfo & getTimezoneInfo() const { return timezone_info; };
-    TimezoneInfo & getTimezoneInfo() { return timezone_info; };
+    const TimezoneInfo & getTimezoneInfo() const { return timezone_info; }
+    TimezoneInfo & getTimezoneInfo() { return timezone_info; }
 
     /// User name and session identifier. Named sessions are local to users.
     using SessionKey = std::pair<String, String>;

dbms/src/Interpreters/Join.cpp

@@ -1385,6 +1385,9 @@ Block Join::doJoinBlockHash(ProbeProcessInfo & probe_process_info, const JoinBui
 
 Block Join::removeUselessColumn(Block & block) const
 {
+    // cancelled
+    if (!block)
+        return block;
     Block projected_block;
     for (const auto & name_and_type : output_columns_after_finalize)
     {
@@ -1415,6 +1418,8 @@ Block Join::joinBlockHash(ProbeProcessInfo & probe_process_info) const
         restore_config.restore_round);
     while (true)
     {
+        if (is_cancelled())
+            return {};
         auto block = doJoinBlockHash(probe_process_info, join_build_info);
         assert(block);
         block = removeUselessColumn(block);
@@ -1519,6 +1524,8 @@ Block Join::joinBlockCross(ProbeProcessInfo & probe_process_info) const
 
     while (true)
     {
+        if (is_cancelled())
+            return {};
         Block block = doJoinBlockCross(probe_process_info);
         assert(block);
         block = removeUselessColumn(block);
@@ -1600,6 +1607,9 @@ Block Join::joinBlockNullAwareSemiImpl(const ProbeProcessInfo & probe_process_in
 
     RUNTIME_ASSERT(res.size() == rows, "SemiJoinResult size {} must be equal to block size {}", res.size(), rows);
 
+    if (is_cancelled())
+        return {};
+
     Block block{};
     for (size_t i = 0; i < probe_process_info.block.columns(); ++i)
     {
@@ -1636,13 +1646,17 @@ Block Join::joinBlockNullAwareSemiImpl(const ProbeProcessInfo & probe_process_in
             blocks,
             null_rows,
             max_block_size,
-            non_equal_conditions);
+            non_equal_conditions,
+            is_cancelled);
 
         helper.joinResult(res_list);
 
         RUNTIME_CHECK_MSG(res_list.empty(), "NASemiJoinResult list must be empty after calculating join result");
     }
 
+    if (is_cancelled())
+        return {};
+
     /// Now all results are known.
 
     std::unique_ptr<IColumn::Filter> filter;
@@ -1787,6 +1801,8 @@ Block Join::joinBlockSemiImpl(const JoinBuildInfo & join_build_info, const Probe
         probe_process_info);
 
     RUNTIME_ASSERT(res.size() == rows, "SemiJoinResult size {} must be equal to block size {}", res.size(), rows);
+    if (is_cancelled())
+        return {};
 
     const NameSet & probe_output_name_set = has_other_condition
         ? output_columns_names_set_for_other_condition_after_finalize
@@ -1821,15 +1837,23 @@ Block Join::joinBlockSemiImpl(const JoinBuildInfo & join_build_info, const Probe
     {
         if (!res_list.empty())
         {
-            SemiJoinHelper<KIND, typename Maps::MappedType>
-                helper(block, left_columns, right_column_indices_to_add, max_block_size, non_equal_conditions);
+            SemiJoinHelper<KIND, typename Maps::MappedType> helper(
+                block,
+                left_columns,
+                right_column_indices_to_add,
+                max_block_size,
+                non_equal_conditions,
+                is_cancelled);
 
             helper.joinResult(res_list);
 
             RUNTIME_CHECK_MSG(res_list.empty(), "SemiJoinResult list must be empty after calculating join result");
         }
     }
 
+    if (is_cancelled())
+        return {};
+
     /// Now all results are known.
 
     std::unique_ptr<IColumn::Filter> filter;
@@ -2216,6 +2240,9 @@ Block Join::joinBlock(ProbeProcessInfo & probe_process_info, bool dry_run) const
     else
         block = joinBlockHash(probe_process_info);
 
+    // if cancelled, just return empty block
+    if (!block)
+        return block;
     /// for (cartesian)antiLeftSemi join, the meaning of "match-helper" is `non-matched` instead of `matched`.
     if (kind == Cross_LeftOuterAnti)
     {
@@ -2488,6 +2515,7 @@ std::optional<RestoreInfo> Join::getOneRestoreStream(size_t max_block_size_)
             restore_join->initBuild(build_sample_block, restore_join_build_concurrency);
             restore_join->setInitActiveBuildThreads();
             restore_join->initProbe(probe_sample_block, restore_join_build_concurrency);
+            restore_join->setCancellationHook(is_cancelled);
             BlockInputStreams restore_scan_hash_map_streams;
             restore_scan_hash_map_streams.resize(restore_join_build_concurrency, nullptr);
             if (needScanHashMapAfterProbe(kind))

dbms/src/Interpreters/Join.h

@@ -25,6 +25,7 @@
 #include <Flash/Coprocessor/JoinInterpreterHelper.h>
 #include <Flash/Coprocessor/RuntimeFilterMgr.h>
 #include <Interpreters/AggregationCommon.h>
+#include <Interpreters/CancellationHook.h>
 #include <Interpreters/ExpressionActions.h>
 #include <Interpreters/HashJoinSpillContext.h>
 #include <Interpreters/JoinHashMap.h>
@@ -311,6 +312,8 @@ class Join
     void flushProbeSideMarkedSpillData(size_t stream_index);
     size_t getProbeCacheColumnThreshold() const { return probe_cache_column_threshold; }
 
+    void setCancellationHook(CancellationHook cancellation_hook) { is_cancelled = cancellation_hook; }
+
     static const String match_helper_prefix;
     static const DataTypePtr match_helper_type;
     static const String flag_mapped_entry_helper_prefix;
@@ -446,6 +449,9 @@ class Join
     // the index of vector is the stream_index.
     std::vector<MarkedSpillData> build_side_marked_spilled_data;
     std::vector<MarkedSpillData> probe_side_marked_spilled_data;
+    CancellationHook is_cancelled{[]() {
+        return false;
+    }};
 
 private:
     /** Set information about structure of right hand of JOIN (joined data).

dbms/src/Interpreters/NullAwareSemiJoinHelper.cpp

@@ -244,13 +244,15 @@ NASemiJoinHelper<KIND, STRICTNESS, Mapped>::NASemiJoinHelper(
     const BlocksList & right_blocks_,
     const std::vector<RowsNotInsertToMap *> & null_rows_,
     size_t max_block_size_,
-    const JoinNonEqualConditions & non_equal_conditions_)
+    const JoinNonEqualConditions & non_equal_conditions_,
+    CancellationHook is_cancelled_)
     : block(block_)
     , left_columns(left_columns_)
     , right_column_indices_to_add(right_column_indices_to_add_)
     , right_blocks(right_blocks_)
     , null_rows(null_rows_)
     , max_block_size(max_block_size_)
+    , is_cancelled(is_cancelled_)
     , non_equal_conditions(non_equal_conditions_)
 {
     static_assert(KIND == NullAware_Anti || KIND == NullAware_LeftOuterAnti || KIND == NullAware_LeftOuterSemi);
@@ -280,17 +282,17 @@ void NASemiJoinHelper<KIND, STRICTNESS, Mapped>::joinResult(std::list<NASemiJoin
         res_list.swap(next_step_res_list);
     }
 
-    if (res_list.empty())
+    if (is_cancelled() || res_list.empty())
         return;
 
     runStep<NASemiJoinStep::NOT_NULL_KEY_CHECK_NULL_ROWS>(res_list, next_step_res_list);
     res_list.swap(next_step_res_list);
-    if (res_list.empty())
+    if (is_cancelled() || res_list.empty())
         return;
 
     runStep<NASemiJoinStep::NULL_KEY_CHECK_NULL_ROWS>(res_list, next_step_res_list);
     res_list.swap(next_step_res_list);
-    if (res_list.empty())
+    if (is_cancelled() || res_list.empty())
         return;
 
     runStepAllBlocks(res_list);
@@ -324,6 +326,8 @@ void NASemiJoinHelper<KIND, STRICTNESS, Mapped>::runStep(
 
     while (!res_list.empty())
     {
+        if (is_cancelled())
+            return;
         MutableColumns columns(block_columns);
         for (size_t i = 0; i < block_columns; ++i)
         {
@@ -384,6 +388,8 @@ void NASemiJoinHelper<KIND, STRICTNESS, Mapped>::runStepAllBlocks(std::list<NASe
         NASemiJoinHelper::Result * res = *res_list.begin();
         for (const auto & right_block : right_blocks)
         {
+            if (is_cancelled())
+                return;
             if (res->getStep() == NASemiJoinStep::DONE)
                 break;