diff --git a/be/src/exec/pipeline/pipeline_fragment_context.cpp b/be/src/exec/pipeline/pipeline_fragment_context.cpp index 2fa064c8a68e09..7ff499fbe2d338 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.cpp +++ b/be/src/exec/pipeline/pipeline_fragment_context.cpp @@ -209,13 +209,16 @@ bool PipelineFragmentContext::notify_close() { // Method like exchange sink buffer will call query ctx cancel. If we add lock here // There maybe dead lock. void PipelineFragmentContext::cancel(const Status reason) { + if (notify_close()) { + return; + } + if (!_cancel_status.update(reason)) { + return; + } LOG_INFO("PipelineFragmentContext::cancel") .tag("query_id", print_id(_query_id)) .tag("fragment_id", _fragment_id) .tag("reason", reason.to_string()); - if (notify_close()) { - return; - } // Timeout is a special error code, we need print current stack to debug timeout issue. if (reason.is()) { auto dbg_str = fmt::format("PipelineFragmentContext is cancelled due to timeout:\n{}", diff --git a/be/src/exec/pipeline/pipeline_fragment_context.h b/be/src/exec/pipeline/pipeline_fragment_context.h index 7243b0214d9fa0..8167e78fba5cde 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.h +++ b/be/src/exec/pipeline/pipeline_fragment_context.h @@ -230,6 +230,8 @@ class PipelineFragmentContext : public TaskExecutionContext { // After prepared, `_total_tasks` is equal to the size of `_tasks`. // When submit fail, `_total_tasks` is equal to the number of tasks submitted. std::atomic _total_tasks = 0; + // The first cancellation reason also gates fragment-level cancellation side effects. + AtomicStatus _cancel_status; std::unique_ptr _fragment_level_profile; // This is used by loading process to report Fragment exec status to FE, FE need fragment status to diff --git a/be/test/exec/pipeline/pipeline_task_test.cpp b/be/test/exec/pipeline/pipeline_task_test.cpp index cd00dff86d9598..4bdb32c66fe89b 100644 --- a/be/test/exec/pipeline/pipeline_task_test.cpp +++ b/be/test/exec/pipeline/pipeline_task_test.cpp @@ -165,6 +165,28 @@ class ThrowStdExceptionTask final : public PipelineTask { std::promise* _close_status; }; +class FragmentCancelLogSink final : public google::LogSink { +public: + void send(google::LogSeverity /*severity*/, const char* /*full_filename*/, + const char* /*base_filename*/, int /*line*/, const google::LogMessageTime& /*time*/, + const char* message, std::size_t message_len) override { + std::string log(message, message_len); + if (log.find("PipelineFragmentContext::cancel") != std::string::npos) { + cancel_count.fetch_add(1, std::memory_order_relaxed); + } + if (log.find("PipelineFragmentContext is cancelled due to timeout") != std::string::npos) { + timeout_dump_count.fetch_add(1, std::memory_order_relaxed); + } + if (log.find("PipelineFragmentContext cancel instance") != std::string::npos) { + instance_cancel_count.fetch_add(1, std::memory_order_relaxed); + } + } + + std::atomic cancel_count {0}; + std::atomic timeout_dump_count {0}; + std::atomic instance_cancel_count {0}; +}; + TEST_F(PipelineTaskTest, TEST_CONSTRUCTOR) { auto num_instances = 1; auto pip_id = 0; @@ -818,6 +840,41 @@ TEST_F(PipelineTaskTest, TEST_SCHEDULER_CATCH_STD_EXCEPTION) { EXPECT_TRUE(_context->is_canceled()); } +TEST_F(PipelineTaskTest, TEST_FRAGMENT_CANCEL_IS_IDEMPOTENT) { + _context->_runtime_state = std::move(_runtime_state); + _context->_total_tasks = 1; + TUniqueId fragment_instance_id; + fragment_instance_id.__set_hi(1); + fragment_instance_id.__set_lo(1); + _context->_fragment_instance_ids.push_back(fragment_instance_id); + + auto* exec_env = ExecEnv::GetInstance(); + bool need_clear_new_load_stream_mgr = exec_env->new_load_stream_mgr() == nullptr; + if (need_clear_new_load_stream_mgr) { + exec_env->set_new_load_stream_mgr(NewLoadStreamMgr::create_unique()); + } + Defer clear_new_load_stream_mgr {[&]() { + if (need_clear_new_load_stream_mgr) { + exec_env->clear_new_load_stream_mgr(); + } + }}; + + FragmentCancelLogSink log_sink; + google::AddLogSink(&log_sink); + Defer remove_log_sink {[&]() { google::RemoveLogSink(&log_sink); }}; + + Status timeout = Status::TimedOut("test timeout"); + _context->cancel(timeout); + _context->cancel(timeout); + _context->cancel(timeout); + _context->cancel(Status::InternalError("later cancellation")); + + EXPECT_EQ(log_sink.cancel_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(log_sink.timeout_dump_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(log_sink.instance_cancel_count.load(std::memory_order_relaxed), 1); + EXPECT_EQ(_context->_cancel_status.status().to_string(), timeout.to_string()); +} + TEST_F(PipelineTaskTest, TEST_FINALIZED_TASK_REJECTS_HYBRID_SUBMIT) { auto num_instances = 1; auto pip_id = 0;