Use one ThreadGroup while pushing to materialized views (and some refactoring for ThreadGroup)

src/Common/CurrentThread.cpp

@@ -90,7 +90,7 @@ void CurrentThread::attachInternalTextLogsQueue(const std::shared_ptr<InternalTe
 }
 
 
-ThreadGroupStatusPtr CurrentThread::getGroup()
+ThreadGroupPtr CurrentThread::getGroup()
 {
     if (unlikely(!current_thread))
         return nullptr;

src/Common/CurrentThread.h

@@ -39,7 +39,7 @@ class CurrentThread
     static ThreadStatus & get();
 
     /// Group to which belongs current thread
-    static ThreadGroupStatusPtr getGroup();
+    static ThreadGroupPtr getGroup();
 
     /// A logs queue used by TCPHandler to pass logs to a client
     static void attachInternalTextLogsQueue(const std::shared_ptr<InternalTextLogsQueue> & logs_queue,
@@ -69,9 +69,9 @@ class CurrentThread
 
     /// You must call one of these methods when create a query child thread:
     /// Add current thread to a group associated with the thread group
-    static void attachToGroup(const ThreadGroupStatusPtr & thread_group);
+    static void attachToGroup(const ThreadGroupPtr & thread_group);
     /// Is useful for a ThreadPool tasks
-    static void attachToGroupIfDetached(const ThreadGroupStatusPtr & thread_group);
+    static void attachToGroupIfDetached(const ThreadGroupPtr & thread_group);
 
     /// Non-master threads call this method in destructor automatically
     static void detachFromGroupIfNotDetached();

src/Common/ThreadStatus.cpp

@@ -61,7 +61,7 @@ static thread_local ThreadStack alt_stack;
 static thread_local bool has_alt_stack = false;
 #endif
 
-ThreadGroupStatus::ThreadGroupStatus()
+ThreadGroup::ThreadGroup()
     : master_thread_id(CurrentThread::get().thread_id)
 {}
 
@@ -119,7 +119,7 @@ ThreadStatus::ThreadStatus()
 #endif
 }
 
-ThreadGroupStatusPtr ThreadStatus::getThreadGroup() const
+ThreadGroupPtr ThreadStatus::getThreadGroup() const
 {
     return thread_group;
 }
@@ -139,7 +139,7 @@ ContextPtr ThreadStatus::getGlobalContext() const
     return global_context.lock();
 }
 
-void ThreadGroupStatus::attachInternalTextLogsQueue(const InternalTextLogsQueuePtr & logs_queue, LogsLevel logs_level)
+void ThreadGroup::attachInternalTextLogsQueue(const InternalTextLogsQueuePtr & logs_queue, LogsLevel logs_level)
 {
     std::lock_guard lock(mutex);
     shared_data.logs_queue_ptr = logs_queue;

src/Common/ThreadStatus.h

@@ -41,7 +41,6 @@ class TaskStatsInfoGetter;
 class InternalTextLogsQueue;
 struct ViewRuntimeData;
 class QueryViewsLog;
-class ThreadGroupSwitcher;
 using InternalTextLogsQueuePtr = std::shared_ptr<InternalTextLogsQueue>;
 using InternalTextLogsQueueWeakPtr = std::weak_ptr<InternalTextLogsQueue>;
 
@@ -58,15 +57,15 @@ using ThreadStatusPtr = ThreadStatus *;
   * Create via CurrentThread::initializeQuery (for queries) or directly (for various background tasks).
   * Use via CurrentThread::getGroup.
   */
-class ThreadGroupStatus;
-using ThreadGroupStatusPtr = std::shared_ptr<ThreadGroupStatus>;
+class ThreadGroup;
+using ThreadGroupPtr = std::shared_ptr<ThreadGroup>;
 
-class ThreadGroupStatus
+class ThreadGroup
 {
 public:
-    ThreadGroupStatus();
+    ThreadGroup();
     using FatalErrorCallback = std::function<void()>;
-    ThreadGroupStatus(ContextPtr query_context_, FatalErrorCallback fatal_error_callback_ = {});
+    ThreadGroup(ContextPtr query_context_, FatalErrorCallback fatal_error_callback_ = {});
 
     /// The first thread created this thread group
     const UInt64 master_thread_id;
@@ -104,9 +103,9 @@ class ThreadGroupStatus
     void attachInternalProfileEventsQueue(const InternalProfileEventsQueuePtr & profile_queue);
 
     /// When new query starts, new thread group is created for it, current thread becomes master thread of the query
-    static ThreadGroupStatusPtr createForQuery(ContextPtr query_context_, FatalErrorCallback fatal_error_callback_ = {});
+    static ThreadGroupPtr createForQuery(ContextPtr query_context_, FatalErrorCallback fatal_error_callback_ = {});
 
-    static ThreadGroupStatusPtr createForBackgroundProcess(ContextPtr storage_context);
+    static ThreadGroupPtr createForBackgroundProcess(ContextPtr storage_context);
 
     std::vector<UInt64> getInvolvedThreadIds() const;
     void linkThread(UInt64 thread_it);
@@ -120,6 +119,21 @@ class ThreadGroupStatus
     std::unordered_set<UInt64> thread_ids;
 };
 
+/**
+ * Since merge is executed with multiple threads, this class
+ * switches the parent MemoryTracker as part of the thread group to account all the memory used.
+ */
+class ThreadGroupSwitcher : private boost::noncopyable
+{
+public:
+    explicit ThreadGroupSwitcher(ThreadGroupPtr thread_group);
+    ~ThreadGroupSwitcher();
+
+private:
+    ThreadGroupPtr prev_thread_group;
+};
+
+
 /**
  * We use **constinit** here to tell the compiler the current_thread variable is initialized.
  * If we didn't help the compiler, then it would most likely add a check before every use of the variable to initialize it if needed.
@@ -163,7 +177,7 @@ class ThreadStatus : public boost::noncopyable
 
 private:
     /// Group of threads, to which this thread attached
-    ThreadGroupStatusPtr thread_group;
+    ThreadGroupPtr thread_group;
 
     /// Is set once
     ContextWeakPtr global_context;
@@ -174,7 +188,7 @@ class ThreadStatus : public boost::noncopyable
     using FatalErrorCallback = std::function<void()>;
     FatalErrorCallback fatal_error_callback;
 
-    ThreadGroupStatus::SharedData local_data;
+    ThreadGroup::SharedData local_data;
 
     bool performance_counters_finalized = false;
 
@@ -215,7 +229,7 @@ class ThreadStatus : public boost::noncopyable
     ThreadStatus();
     ~ThreadStatus();
 
-    ThreadGroupStatusPtr getThreadGroup() const;
+    ThreadGroupPtr getThreadGroup() const;
 
     const String & getQueryId() const;
 
@@ -239,7 +253,7 @@ class ThreadStatus : public boost::noncopyable
     void setInternalThread();
 
     /// Attaches slave thread to existing thread group
-    void attachToGroup(const ThreadGroupStatusPtr & thread_group_, bool check_detached = true);
+    void attachToGroup(const ThreadGroupPtr & thread_group_, bool check_detached = true);
 
     /// Detaches thread from the thread group and the query, dumps performance counters if they have not been dumped
     void detachFromGroup();
@@ -287,7 +301,7 @@ class ThreadStatus : public boost::noncopyable
 
     void logToQueryThreadLog(QueryThreadLog & thread_log, const String & current_database);
 
-    void attachToGroupImpl(const ThreadGroupStatusPtr & thread_group_);
+    void attachToGroupImpl(const ThreadGroupPtr & thread_group_);
 };
 
 /**

src/Interpreters/Aggregator.cpp

@@ -2315,7 +2315,7 @@ BlocksList Aggregator::prepareBlocksAndFillTwoLevelImpl(
 
     std::atomic<UInt32> next_bucket_to_merge = 0;
 
-    auto converter = [&](size_t thread_id, ThreadGroupStatusPtr thread_group)
+    auto converter = [&](size_t thread_id, ThreadGroupPtr thread_group)
     {
         SCOPE_EXIT_SAFE(
             if (thread_group)
@@ -3043,7 +3043,7 @@ void Aggregator::mergeBlocks(BucketToBlocks bucket_to_blocks, AggregatedDataVari
 
         LOG_TRACE(log, "Merging partially aggregated two-level data.");
 
-        auto merge_bucket = [&bucket_to_blocks, &result, this](Int32 bucket, Arena * aggregates_pool, ThreadGroupStatusPtr thread_group)
+        auto merge_bucket = [&bucket_to_blocks, &result, this](Int32 bucket, Arena * aggregates_pool, ThreadGroupPtr thread_group)
         {
             SCOPE_EXIT_SAFE(
                 if (thread_group)

src/Interpreters/ExternalLoader.cpp

@@ -967,7 +967,7 @@ class ExternalLoader::LoadingDispatcher : private boost::noncopyable
     }
 
     /// Does the loading, possibly in the separate thread.
-    void doLoading(const String & name, size_t loading_id, bool forced_to_reload, size_t min_id_to_finish_loading_dependencies_, bool async, ThreadGroupStatusPtr thread_group = {})
+    void doLoading(const String & name, size_t loading_id, bool forced_to_reload, size_t min_id_to_finish_loading_dependencies_, bool async, ThreadGroupPtr thread_group = {})
     {
         SCOPE_EXIT_SAFE(
             if (thread_group)

src/Interpreters/InterpreterInsertQuery.cpp

@@ -32,6 +32,7 @@
 #include <Storages/StorageMaterializedView.h>
 #include <Storages/WindowView/StorageWindowView.h>
 #include <TableFunctions/TableFunctionFactory.h>
+#include <Common/ThreadStatus.h>
 #include <Common/checkStackSize.h>
 
 
@@ -233,15 +234,22 @@ Chain InterpreterInsertQuery::buildChain(
     ThreadStatusesHolderPtr thread_status_holder,
     std::atomic_uint64_t * elapsed_counter_ms)
 {
+    ThreadGroupPtr running_group;
+    if (current_thread)
+        running_group = current_thread->getThreadGroup();
+    if (!running_group)
+        running_group = std::make_shared<ThreadGroup>(getContext());
+
     auto sample = getSampleBlock(columns, table, metadata_snapshot);
-    return buildChainImpl(table, metadata_snapshot, sample, thread_status_holder, elapsed_counter_ms);
+    return buildChainImpl(table, metadata_snapshot, sample, thread_status_holder, running_group, elapsed_counter_ms);
 }
 
 Chain InterpreterInsertQuery::buildChainImpl(
     const StoragePtr & table,
     const StorageMetadataPtr & metadata_snapshot,
     const Block & query_sample_block,
     ThreadStatusesHolderPtr thread_status_holder,
+    ThreadGroupPtr running_group,
     std::atomic_uint64_t * elapsed_counter_ms)
 {
     ThreadStatus * thread_status = current_thread;
@@ -273,7 +281,9 @@ Chain InterpreterInsertQuery::buildChainImpl(
     }
     else
     {
-        out = buildPushingToViewsChain(table, metadata_snapshot, context_ptr, query_ptr, no_destination, thread_status_holder, elapsed_counter_ms);
+        out = buildPushingToViewsChain(table, metadata_snapshot, context_ptr,
+            query_ptr, no_destination,
+            thread_status_holder, running_group, elapsed_counter_ms);
     }
 
     /// Note that we wrap transforms one on top of another, so we write them in reverse of data processing order.
@@ -461,9 +471,17 @@ BlockIO InterpreterInsertQuery::execute()
             pipeline = interpreter_watch.buildQueryPipeline();
         }
 
+        ThreadGroupPtr running_group;
+        if (current_thread)
+            running_group = current_thread->getThreadGroup();
+        if (!running_group)
+            running_group = std::make_shared<ThreadGroup>(getContext());
         for (size_t i = 0; i < out_streams_size; ++i)
         {
-            auto out = buildChainImpl(table, metadata_snapshot, query_sample_block, nullptr, nullptr);
+            auto out = buildChainImpl(table, metadata_snapshot, query_sample_block,
+                /* thread_status_holder= */ nullptr,
+                running_group,
+                /* elapsed_counter_ms= */ nullptr);
             out_chains.emplace_back(std::move(out));
         }
     }

src/Interpreters/InterpreterInsertQuery.h

@@ -70,6 +70,7 @@ class InterpreterInsertQuery : public IInterpreter, WithContext
         const StorageMetadataPtr & metadata_snapshot,
         const Block & query_sample_block,
         ThreadStatusesHolderPtr thread_status_holder,
+        ThreadGroupPtr running_group,
         std::atomic_uint64_t * elapsed_counter_ms);
 };
 

src/Interpreters/ProcessList.cpp

@@ -340,7 +340,7 @@ QueryStatus::QueryStatus(
     const String & query_,
     const ClientInfo & client_info_,
     QueryPriorities::Handle && priority_handle_,
-    ThreadGroupStatusPtr && thread_group_,
+    ThreadGroupPtr && thread_group_,
     IAST::QueryKind query_kind_,
     UInt64 watch_start_nanoseconds)
     : WithContext(context_)

src/Interpreters/ProcessList.h

@@ -86,7 +86,7 @@ class QueryStatus : public WithContext
     ClientInfo client_info;
 
     /// Info about all threads involved in query execution
-    ThreadGroupStatusPtr thread_group;
+    ThreadGroupPtr thread_group;
 
     Stopwatch watch;
 
@@ -162,7 +162,7 @@ class QueryStatus : public WithContext
         const String & query_,
         const ClientInfo & client_info_,
         QueryPriorities::Handle && priority_handle_,
-        ThreadGroupStatusPtr && thread_group_,
+        ThreadGroupPtr && thread_group_,
         IAST::QueryKind query_kind_,
         UInt64 watch_start_nanoseconds);