join be aware of cancel signal (#9450)

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

@@ -520,6 +520,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

@@ -20,6 +20,7 @@
 #include <Debug/MockServerInfo.h>
 #include <Encryption/FileProvider_fwd.h>
 #include <IO/CompressionSettings.h>
+#include <Interpreters/CancellationHook.h>
 #include <Interpreters/ClientInfo.h>
 #include <Interpreters/Context_fwd.h>
 #include <Interpreters/Settings.h>
@@ -181,6 +182,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
@@ -238,8 +242,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,
@@ -376,6 +380,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;
@@ -502,8 +509,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

@@ -1352,6 +1352,10 @@ Block Join::doJoinBlockHash(ProbeProcessInfo & probe_process_info) const
 
 Block Join::removeUselessColumn(Block & block) const
 {
+    // cancelled
+    if (!block)
+        return block;
+
     Block projected_block;
     for (const auto & name : tidb_output_column_names)
     {
@@ -1368,6 +1372,8 @@ Block Join::joinBlockHash(ProbeProcessInfo & probe_process_info) const
     probe_process_info.prepareForHashProbe(key_names_left, non_equal_conditions.left_filter_column, kind, strictness);
     while (true)
     {
+        if (is_cancelled())
+            return {};
         auto block = doJoinBlockHash(probe_process_info);
         assert(block);
         block = removeUselessColumn(block);
@@ -1468,6 +1474,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);
@@ -1582,6 +1590,8 @@ Block Join::joinBlockNullAware(ProbeProcessInfo & probe_process_info) const
 
     /// Null aware join never expand the left block, just handle the whole block at one time is enough
     probe_process_info.all_rows_joined_finish = true;
+    if (is_cancelled())
+        return {};
 
     return removeUselessColumn(block);
 }
@@ -1616,19 +1626,31 @@ void Join::joinBlockNullAwareImpl(
         right_side_info);
 
     RUNTIME_ASSERT(res.size() == rows, "NASemiJoinResult size {} must be equal to block size {}", res.size(), rows);
+    if (is_cancelled())
+        return;
 
     size_t right_columns = block.columns() - left_columns;
 
     if (!res_list.empty())
     {
-        NASemiJoinHelper<KIND, STRICTNESS, typename Maps::MappedType::Base_t>
-            helper(block, left_columns, right_columns, blocks, null_rows, max_block_size, non_equal_conditions);
+        NASemiJoinHelper<KIND, STRICTNESS, typename Maps::MappedType::Base_t> helper(
+            block,
+            left_columns,
+            right_columns,
+            blocks,
+            null_rows,
+            max_block_size,
+            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;
@@ -2004,6 +2026,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 == LeftOuterAnti || kind == Cross_LeftOuterAnti)
     {
@@ -2275,6 +2300,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>
@@ -306,6 +307,8 @@ class Join
     bool hasProbeSideMarkedSpillData(size_t stream_index) const;
     void flushProbeSideMarkedSpillData(size_t stream_index);
 
+    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;
@@ -433,6 +436,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

@@ -230,13 +230,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_columns(right_columns_)
     , 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);
@@ -265,17 +267,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);
@@ -309,6 +311,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)
         {
@@ -369,6 +373,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;
 

dbms/src/Interpreters/NullAwareSemiJoinHelper.h

@@ -17,6 +17,7 @@
 #include <Common/Logger.h>
 #include <Core/Block.h>
 #include <Flash/Coprocessor/JoinInterpreterHelper.h>
+#include <Interpreters/CancellationHook.h>
 #include <Parsers/ASTTablesInSelectQuery.h>
 
 
@@ -173,7 +174,8 @@ class 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_);
 
     void joinResult(std::list<Result *> & res_list);
 
@@ -197,6 +199,7 @@ class NASemiJoinHelper
     const BlocksList & right_blocks;
     const std::vector<RowsNotInsertToMap *> & null_rows;
     size_t max_block_size;
+    CancellationHook is_cancelled;
 
     const JoinNonEqualConditions & non_equal_conditions;
 };

dbms/src/Operators/AggregateContext.cpp

@@ -19,7 +19,7 @@ namespace DB
 void AggregateContext::initBuild(
     const Aggregator::Params & params,
     size_t max_threads_,
-    Aggregator::CancellationHook && hook,
+    CancellationHook && hook,
     const RegisterOperatorSpillContext & register_operator_spill_context)
 {
     assert(status.load() == AggStatus::init);