From 333bba4e6b6ed30e54544273586fc4ebd6a51c01 Mon Sep 17 00:00:00 2001
From: Yarik Briukhovetskyi <114298166+yariks5s@users.noreply.github.com>
Date: Tue, 19 Nov 2024 12:35:42 +0100
Subject: [PATCH] Add locks on exception throwing.

---
 src/Interpreters/ProcessList.cpp | 10 ++++++----
 src/Interpreters/ProcessList.h   |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/Interpreters/ProcessList.cpp b/src/Interpreters/ProcessList.cpp
index 0ab4d292e39..08a54f85302 100644
--- a/src/Interpreters/ProcessList.cpp
+++ b/src/Interpreters/ProcessList.cpp
@@ -459,7 +459,6 @@ void QueryStatus::ExecutorHolder::remove()
 CancellationCode QueryStatus::cancelQuery(CancelReason reason, std::exception_ptr exception)
 {
     {
-        // Lock the mutex to protect the critical section
         std::lock_guard<std::mutex> lock(cancel_mutex);
 
         if (is_killed)
@@ -501,9 +500,12 @@ CancellationCode QueryStatus::cancelQuery(CancelReason reason, std::exception_pt
 
 void QueryStatus::throwProperExceptionIfNeeded(const UInt64 & max_execution_time, const UInt64 & elapsed_ns) const
 {
-    if (is_killed.load())
     {
-        throwProperException(max_execution_time, elapsed_ns);
+        std::lock_guard<std::mutex> lock(cancel_mutex);
+        if (is_killed)
+        {
+            throwProperException(max_execution_time, elapsed_ns);
+        }
     }
 }
 
@@ -515,7 +517,7 @@ void QueryStatus::throwProperException(const UInt64 & max_execution_time, const
 
     if (cancel_reason == CancelReason::TIMEOUT)
         throw Exception(ErrorCodes::TIMEOUT_EXCEEDED, "Timeout exceeded: {} maximum: {} ms", additional_error_part, max_execution_time / 1000.0);
-    throw Exception(ErrorCodes::QUERY_WAS_CANCELLED, "Query was cancelled");
+    throwQueryWasCancelled();
 }
 
 void QueryStatus::addPipelineExecutor(PipelineExecutor * e)
diff --git a/src/Interpreters/ProcessList.h b/src/Interpreters/ProcessList.h
index ba7e1ec00af..7ee6008d79c 100644
--- a/src/Interpreters/ProcessList.h
+++ b/src/Interpreters/ProcessList.h
@@ -116,7 +116,7 @@ protected:
     bool is_cancelling { false };
     /// KILL was send to the query
     std::atomic<bool> is_killed { false };
-    std::atomic<CancelReason> cancel_reason { CancelReason::UNDEFINED };
+    CancelReason cancel_reason { CancelReason::UNDEFINED };
 
     std::exception_ptr cancellation_exception;
     mutable std::mutex cancellation_exception_mutex;