Set query_id for mutations/merges

This will allow to distinguish allocations in trace_log.

src/Common/ThreadStatus.h

@@ -41,6 +41,7 @@ class TaskStatsInfoGetter;
 class InternalTextLogsQueue;
 struct ViewRuntimeData;
 class QueryViewsLog;
+class MemoryTrackerThreadSwitcher;
 using InternalTextLogsQueuePtr = std::shared_ptr<InternalTextLogsQueue>;
 using InternalTextLogsQueueWeakPtr = std::weak_ptr<InternalTextLogsQueue>;
 
@@ -166,6 +167,13 @@ protected:
     /// It is used to avoid enabling the query profiler when you have multiple ThreadStatus in the same thread
     bool query_profiled_enabled = true;
 
+    /// Requires access to query_id.
+    friend class MemoryTrackerThreadSwitcher;
+    void setQueryId(const String & query_id_)
+    {
+        query_id = query_id_;
+    }
+
 public:
     ThreadStatus();
     ~ThreadStatus();

src/Storages/MergeTree/MergeList.cpp

@@ -10,7 +10,7 @@ namespace DB
 {
 
 
-MemoryTrackerThreadSwitcher::MemoryTrackerThreadSwitcher(MemoryTracker * memory_tracker_ptr, UInt64 untracked_memory_limit)
+MemoryTrackerThreadSwitcher::MemoryTrackerThreadSwitcher(MemoryTracker * memory_tracker_ptr, UInt64 untracked_memory_limit, const std::string & query_id)
 {
     // Each merge is executed into separate background processing pool thread
     background_thread_memory_tracker = CurrentThread::getMemoryTracker();
@@ -39,6 +39,9 @@ MemoryTrackerThreadSwitcher::MemoryTrackerThreadSwitcher(MemoryTracker * memory_
     /// (NOTE: consider moving such code to ThreadFromGlobalPool and related places)
     prev_untracked_memory = current_thread->untracked_memory;
     current_thread->untracked_memory = 0;
+
+    prev_query_id = current_thread->getQueryId().toString();
+    current_thread->setQueryId(query_id);
 }
 
 
@@ -51,6 +54,7 @@ MemoryTrackerThreadSwitcher::~MemoryTrackerThreadSwitcher()
 
     current_thread->untracked_memory_limit = prev_untracked_memory_limit;
     current_thread->untracked_memory = prev_untracked_memory;
+    current_thread->setQueryId(prev_query_id);
 }
 
 MergeListElement::MergeListElement(
@@ -66,6 +70,7 @@ MergeListElement::MergeListElement(
     , result_part_info{future_part->part_info}
     , num_parts{future_part->parts.size()}
     , max_untracked_memory(max_untracked_memory_)
+    , query_id(table_id.getShortName() + "::" + result_part_name)
     , thread_id{getThreadId()}
     , merge_type{future_part->merge_type}
     , merge_algorithm{MergeAlgorithm::Undecided}

src/Storages/MergeTree/MergeList.h

@@ -62,13 +62,14 @@ using FutureMergedMutatedPartPtr = std::shared_ptr<FutureMergedMutatedPart>;
 class MemoryTrackerThreadSwitcher : boost::noncopyable
 {
 public:
-    explicit MemoryTrackerThreadSwitcher(MemoryTracker * memory_tracker_ptr, UInt64 untracked_memory_limit);
+    explicit MemoryTrackerThreadSwitcher(MemoryTracker * memory_tracker_ptr, UInt64 untracked_memory_limit, const std::string & query_id);
     ~MemoryTrackerThreadSwitcher();
 private:
     MemoryTracker * background_thread_memory_tracker;
     MemoryTracker * background_thread_memory_tracker_prev_parent = nullptr;
     UInt64 prev_untracked_memory_limit;
     UInt64 prev_untracked_memory;
+    String prev_query_id;
 };
 
 using MemoryTrackerThreadSwitcherPtr = std::unique_ptr<MemoryTrackerThreadSwitcher>;
@@ -107,6 +108,7 @@ struct MergeListElement : boost::noncopyable
 
     MemoryTracker memory_tracker{VariableContext::Process};
     UInt64 max_untracked_memory;
+    std::string query_id;
 
     UInt64 thread_id;
     MergeType merge_type;

src/Storages/MergeTree/MergePlainMergeTreeTask.cpp

@@ -30,7 +30,10 @@ bool MergePlainMergeTreeTask::executeStep()
     /// Make out memory tracker a parent of current thread memory tracker
     MemoryTrackerThreadSwitcherPtr switcher;
     if (merge_list_entry)
-        switcher = std::make_unique<MemoryTrackerThreadSwitcher>(&(*merge_list_entry)->memory_tracker, (*merge_list_entry)->max_untracked_memory);
+        switcher = std::make_unique<MemoryTrackerThreadSwitcher>(
+            &(*merge_list_entry)->memory_tracker,
+            (*merge_list_entry)->max_untracked_memory,
+            (*merge_list_entry)->query_id);
 
     switch (state)
     {

src/Storages/MergeTree/MutatePlainMergeTreeTask.cpp

@@ -62,7 +62,10 @@ bool MutatePlainMergeTreeTask::executeStep()
     /// Make out memory tracker a parent of current thread memory tracker
     MemoryTrackerThreadSwitcherPtr switcher;
     if (merge_list_entry)
-        switcher = std::make_unique<MemoryTrackerThreadSwitcher>(&(*merge_list_entry)->memory_tracker, (*merge_list_entry)->max_untracked_memory);
+        switcher = std::make_unique<MemoryTrackerThreadSwitcher>(
+            &(*merge_list_entry)->memory_tracker,
+            (*merge_list_entry)->max_untracked_memory,
+            (*merge_list_entry)->query_id);
 
     switch (state)
     {

src/Storages/MergeTree/ReplicatedMergeMutateTaskBase.cpp

@@ -122,7 +122,10 @@ bool ReplicatedMergeMutateTaskBase::executeImpl()
 {
     MemoryTrackerThreadSwitcherPtr switcher;
     if (merge_mutate_entry)
-        switcher = std::make_unique<MemoryTrackerThreadSwitcher>(&(*merge_mutate_entry)->memory_tracker, (*merge_mutate_entry)->max_untracked_memory);
+        switcher = std::make_unique<MemoryTrackerThreadSwitcher>(
+            &(*merge_mutate_entry)->memory_tracker,
+            (*merge_mutate_entry)->max_untracked_memory,
+            (*merge_mutate_entry)->query_id);
 
     auto remove_processed_entry = [&] () -> bool
     {