[fix](pipeline) Make fragment cancellation idempotent

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<ErrorCode::TIMEOUT>()) {
         auto dbg_str = fmt::format("PipelineFragmentContext is cancelled due to timeout:\n{}",

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<int> _total_tasks = 0;
+    // The first cancellation reason also gates fragment-level cancellation side effects.
+    AtomicStatus _cancel_status;
 
     std::unique_ptr<RuntimeProfile> _fragment_level_profile;
     // This is used by loading process to report Fragment exec status to FE, FE need fragment status to

be/test/exec/pipeline/pipeline_task_test.cpp

@@ -165,6 +165,28 @@ class ThrowStdExceptionTask final : public PipelineTask {
     std::promise<std::string>* _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<int> cancel_count {0};
+    std::atomic<int> timeout_dump_count {0};
+    std::atomic<int> 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;