From 7599020b8758f4e835a57d30d48b868a72182434 Mon Sep 17 00:00:00 2001
From: Yarik Briukhovetskyi <114298166+yariks5s@users.noreply.github.com>
Date: Tue, 19 Nov 2024 20:14:23 +0100
Subject: [PATCH] Fix the issue in the comment above.

---
 src/Interpreters/CancellationChecker.cpp |  4 ++--
 src/Interpreters/CancellationChecker.h   |  2 +-
 src/Interpreters/ProcessList.cpp         | 24 +++++++++++-------------
 src/Interpreters/ProcessList.h           |  3 +--
 4 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/src/Interpreters/CancellationChecker.cpp b/src/Interpreters/CancellationChecker.cpp
index 4554bea301b..e0c1f88f89d 100644
--- a/src/Interpreters/CancellationChecker.cpp
+++ b/src/Interpreters/CancellationChecker.cpp
@@ -42,7 +42,7 @@ void CancellationChecker::terminateThread()
     cond_var.notify_all();
 }
 
-void CancellationChecker::cancelTask(QueryToTrack task, [[maybe_unused]]CancelReason reason)
+void CancellationChecker::cancelTask(QueryToTrack task)
 {
     if (task.query)
     {
@@ -123,7 +123,7 @@ void CancellationChecker::workerFunction()
             if ((end_time_ms <= now_ms && duration_milliseconds.count() != 0))
             {
                 LOG_TRACE(getLogger("CancellationChecker"), "Cancelling the task because of the timeout: {} ms, query: {}", duration, next_task.query->getInfo().query);
-                cancelTask(next_task, CancelReason::TIMEOUT);
+                cancelTask(next_task);
                 querySet.erase(next_task);
 
                 continue;
diff --git a/src/Interpreters/CancellationChecker.h b/src/Interpreters/CancellationChecker.h
index e890fc9aec3..3addc2a0cc2 100644
--- a/src/Interpreters/CancellationChecker.h
+++ b/src/Interpreters/CancellationChecker.h
@@ -46,7 +46,7 @@ private:
     std::condition_variable cond_var;
 
     // Function to execute when a task's endTime is reached
-    void cancelTask(QueryToTrack task, CancelReason reason);
+    void cancelTask(QueryToTrack task);
     bool removeQueryFromSet(std::shared_ptr<QueryStatus> query);
 
 public:
diff --git a/src/Interpreters/ProcessList.cpp b/src/Interpreters/ProcessList.cpp
index c2c280badc8..f2382177ddf 100644
--- a/src/Interpreters/ProcessList.cpp
+++ b/src/Interpreters/ProcessList.cpp
@@ -498,28 +498,26 @@ CancellationCode QueryStatus::cancelQuery(CancelReason reason, std::exception_pt
     return CancellationCode::CancelSent;
 }
 
-void QueryStatus::throwProperExceptionIfNeeded(const UInt64 & max_execution_time, const UInt64 & elapsed_ns) const
+void QueryStatus::throwProperExceptionIfNeeded(const UInt64 & max_execution_time, const UInt64 & elapsed_ns)
 {
     {
         std::lock_guard<std::mutex> lock(cancel_mutex);
         if (is_killed)
         {
-            throwProperException(max_execution_time, elapsed_ns);
+            String additional_error_part;
+            if (!elapsed_ns)
+                additional_error_part = fmt::format("elapsed {} ms, ", static_cast<double>(elapsed_ns) / 1000000000ULL);
+
+            if (cancel_reason == CancelReason::TIMEOUT)
+            {
+                cancel_reason = CancelReason::UNDEFINED; // We can assign only CancelReason::TIMEOUT to cancel_reason, so we need to assign it back to UNDEFINED
+                throw Exception(ErrorCodes::TIMEOUT_EXCEEDED, "Timeout exceeded: {}maximum: {} ms", additional_error_part, max_execution_time / 1000.0);
+            }
+            throwQueryWasCancelled();
         }
     }
 }
 
-void QueryStatus::throwProperException(const UInt64 & max_execution_time, const UInt64 & elapsed_ns) const
-{
-    String additional_error_part;
-    if (!elapsed_ns)
-        additional_error_part = fmt::format("elapsed {} ms,", static_cast<double>(elapsed_ns) / 1000000000ULL);
-
-    if (cancel_reason == CancelReason::TIMEOUT)
-        throw Exception(ErrorCodes::TIMEOUT_EXCEEDED, "Timeout exceeded: {} maximum: {} ms", additional_error_part, max_execution_time / 1000.0);
-    throwQueryWasCancelled();
-}
-
 void QueryStatus::addPipelineExecutor(PipelineExecutor * e)
 {
     /// In case of asynchronous distributed queries it is possible to call
diff --git a/src/Interpreters/ProcessList.h b/src/Interpreters/ProcessList.h
index a5effc9cb51..1d4bcf8b758 100644
--- a/src/Interpreters/ProcessList.h
+++ b/src/Interpreters/ProcessList.h
@@ -240,8 +240,7 @@ public:
 
     QueryStatusInfo getInfo(bool get_thread_list = false, bool get_profile_events = false, bool get_settings = false) const;
 
-    void throwProperExceptionIfNeeded(const UInt64 & max_execution_time, const UInt64 & elapsed_ns = 0) const;
-    [[noreturn]] void throwProperException(const UInt64 & max_execution_time, const UInt64 & elapsed_ns = 0) const;
+    void throwProperExceptionIfNeeded(const UInt64 & max_execution_time, const UInt64 & elapsed_ns = 0);
 
     /// Cancels the current query.
     /// Optional argument `exception` allows to set an exception which checkTimeLimit() will throw instead of "QUERY_WAS_CANCELLED".