24.8.14 Backport of #75474 - Stricter checks when attaching/detaching threads to groups

src/AggregateFunctions/UniqExactSet.h

@@ -65,14 +65,7 @@ class UniqExactSet
                 auto data_vec_atomic_index = std::make_shared<std::atomic_uint32_t>(0);
                 auto thread_func = [data_vec, data_vec_atomic_index, &is_cancelled, thread_group = CurrentThread::getGroup()]()
                 {
-                    SCOPE_EXIT_SAFE(
-                        if (thread_group)
-                            CurrentThread::detachFromGroupIfNotDetached();
-                    );
-                    if (thread_group)
-                        CurrentThread::attachToGroupIfDetached(thread_group);
-
-                    setThreadName("UniqExaConvert");
+                    ThreadGroupSwitcher switcher(thread_group, "UniqExaConvert");
 
                     while (true)
                     {
@@ -128,13 +121,7 @@ class UniqExactSet
 
                     auto thread_func = [&lhs, &rhs, next_bucket_to_merge, is_cancelled, thread_group = CurrentThread::getGroup()]()
                     {
-                        SCOPE_EXIT_SAFE(
-                            if (thread_group)
-                                CurrentThread::detachFromGroupIfNotDetached();
-                        );
-                        if (thread_group)
-                            CurrentThread::attachToGroupIfDetached(thread_group);
-                        setThreadName("UniqExactMerger");
+                        ThreadGroupSwitcher switcher(thread_group, "UniqExactMerger");
 
                         while (true)
                         {

src/Common/CurrentMetrics.cpp

@@ -228,6 +228,9 @@
     M(OutdatedPartsLoadingThreads, "Number of threads in the threadpool for loading Outdated data parts.") \
     M(OutdatedPartsLoadingThreadsActive, "Number of active threads in the threadpool for loading Outdated data parts.") \
     M(OutdatedPartsLoadingThreadsScheduled, "Number of queued or active jobs in the threadpool for loading Outdated data parts.") \
+    M(PolygonDictionaryThreads, "Number of threads in the threadpool for polygon dictionaries.") \
+    M(PolygonDictionaryThreadsActive, "Number of active threads in the threadpool for polygon dictionaries.") \
+    M(PolygonDictionaryThreadsScheduled, "Number of queued or active jobs in the threadpool for polygon dictionaries.") \
     M(DistributedBytesToInsert, "Number of pending bytes to process for asynchronous insertion into Distributed tables. Number of bytes for every shard is summed.") \
     M(BrokenDistributedBytesToInsert, "Number of bytes for asynchronous insertion into Distributed tables that has been marked as broken. Number of bytes for every shard is summed.") \
     M(DistributedFilesToInsert, "Number of pending files to process for asynchronous insertion into Distributed tables. Number of files for every shard is summed.") \

src/Common/ThreadPool.cpp

@@ -1,4 +1,5 @@
 #include <Common/ThreadPool.h>
+#include <Common/CurrentThread.h>
 #include <Common/ProfileEvents.h>
 #include <Common/setThreadName.h>
 #include <Common/Exception.h>
@@ -769,6 +770,11 @@ void ThreadPoolImpl<Thread>::ThreadFromThreadPool::worker()
 
             CurrentMetrics::Increment metric_active_pool_threads(parent_pool.metric_active_threads);
 
+#ifdef DEBUG_OR_SANITIZER_BUILD
+            DB::ThreadStatus * initial_thread = DB::current_thread;
+            DB::ThreadGroupPtr initial_thread_group = DB::CurrentThread::getGroup();
+#endif
+
             if constexpr (!std::is_same_v<Thread, std::thread>)
             {
                 Stopwatch watch;
@@ -785,6 +791,13 @@ void ThreadPoolImpl<Thread>::ThreadFromThreadPool::worker()
                 job_data->job();
             }
 
+#ifdef DEBUG_OR_SANITIZER_BUILD
+            DB::ThreadStatus * final_thread = DB::current_thread;
+            DB::ThreadGroupPtr final_thread_group = DB::CurrentThread::getGroup();
+            if (final_thread != initial_thread || final_thread_group != initial_thread_group)
+                throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Thread pool job changed current ThreadStatus pointer ({} -> {}) or ThreadGroup ({} -> {}).", initial_thread ? "non-nullptr" : "nullptr", final_thread ? "non-nullptr" : "nullptr", initial_thread_group ? "master_thread_id " + std::to_string(initial_thread_group->master_thread_id) : "nullptr", final_thread_group ? "master_thread_id " + std::to_string(final_thread_group->master_thread_id) : "nullptr");
+#endif
+
 
             if (thread_trace_context.root_span.isTraceEnabled())
             {

src/Common/ThreadStatus.h

@@ -133,17 +133,31 @@ class ThreadGroup
 };
 
 /**
- * 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.
+ * RAII wrapper around CurrentThread::attachToGroup/detachFromGroupIfNotDetached.
+ *
+ * Typically used for inheriting thread group when scheduling tasks on a thread pool:
+ *   pool->scheduleOrThrow([thread_group = CurrentThread::getGroup()]()
+ *       {
+ *           ThreadGroupSwitcher switcher(thread_group, "MyThread");
+ *           ...
+ *       });
  */
 class ThreadGroupSwitcher : private boost::noncopyable
 {
 public:
-    explicit ThreadGroupSwitcher(ThreadGroupPtr thread_group);
+    /// If thread_group_ is nullptr or equal to current thread group, does nothing.
+    /// allow_existing_group:
+    ///  * If false, asserts that the thread is not already attached to a different group.
+    ///    Use this when running a task in a thread pool.
+    ///  * If true, remembers the current group and restores it in destructor.
+    /// If thread_name is not empty, calls setThreadName along the way; should be at most 15 bytes long.
+    explicit ThreadGroupSwitcher(ThreadGroupPtr thread_group_, const char * thread_name, bool allow_existing_group = false) noexcept;
     ~ThreadGroupSwitcher();
 
 private:
+    ThreadStatus * prev_thread = nullptr;
     ThreadGroupPtr prev_thread_group;
+    ThreadGroupPtr thread_group;
 };
 
 

src/Common/threadPoolCallbackRunner.h

@@ -31,24 +31,15 @@ ThreadPoolCallbackRunnerUnsafe<Result, Callback> threadPoolCallbackRunnerUnsafe(
     {
         auto task = std::make_shared<std::packaged_task<Result()>>([thread_group, thread_name, my_callback = std::move(callback)]() mutable -> Result
         {
-            if (thread_group)
-                CurrentThread::attachToGroup(thread_group);
+            ThreadGroupSwitcher switcher(thread_group, thread_name.c_str());
 
             SCOPE_EXIT_SAFE(
             {
-                {
-                    /// Release all captured resources before detaching thread group
-                    /// Releasing has to use proper memory tracker which has been set here before callback
-
-                    [[maybe_unused]] auto tmp = std::move(my_callback);
-                }
-
-                if (thread_group)
-                    CurrentThread::detachFromGroupIfNotDetached();
+                /// Release all captured resources before detaching thread group
+                /// Releasing has to use proper memory tracker which has been set here before callback
+                [[maybe_unused]] auto tmp = std::move(my_callback);
             });
 
-            setThreadName(thread_name.data());
-
             return my_callback();
         });
 
@@ -75,6 +66,8 @@ std::future<Result> scheduleFromThreadPoolUnsafe(T && task, ThreadPool & pool, c
 template <typename Result, typename PoolT = ThreadPool, typename Callback = std::function<Result()>>
 class ThreadPoolCallbackRunnerLocal
 {
+    static_assert(!std::is_same_v<PoolT, GlobalThreadPool>, "Scheduling tasks directly on GlobalThreadPool is not allowed because it doesn't set up CurrentThread. Create a new ThreadPool (local or in SharedThreadPools.h) or use ThreadFromGlobalPool.");
+
     PoolT & pool;
     std::string thread_name;
 
@@ -124,6 +117,8 @@ class ThreadPoolCallbackRunnerLocal
         auto task_func = std::make_shared<std::packaged_task<Result()>>(
         [task, thread_group = CurrentThread::getGroup(), my_thread_name = thread_name, my_callback = std::move(callback)]() mutable -> Result
         {
+            ThreadGroupSwitcher switcher(thread_group, my_thread_name.c_str());
+
             TaskState expected = SCHEDULED;
             if (!task->state.compare_exchange_strong(expected, RUNNING))
             {
@@ -139,9 +134,6 @@ class ThreadPoolCallbackRunnerLocal
                     throw Exception(ErrorCodes::LOGICAL_ERROR, "Unexpected state {} when finishing a task in {}", expected, my_thread_name);
             });
 
-            if (thread_group)
-                CurrentThread::attachToGroup(thread_group);
-
             SCOPE_EXIT_SAFE(
             {
                 {
@@ -150,13 +142,8 @@ class ThreadPoolCallbackRunnerLocal
 
                     [[maybe_unused]] auto tmp = std::move(my_callback);
                 }
-
-                if (thread_group)
-                    CurrentThread::detachFromGroupIfNotDetached();
             });
 
-            setThreadName(my_thread_name.data());
-
             return my_callback();
         });
 

src/Dictionaries/HashedDictionary.h

@@ -336,18 +336,11 @@ HashedDictionary<dictionary_key_type, sparse, sharded>::~HashedDictionary()
 
         pool.trySchedule([&container, thread_group = CurrentThread::getGroup()]
         {
-            SCOPE_EXIT_SAFE(
-                if (thread_group)
-                    CurrentThread::detachFromGroupIfNotDetached();
-            );
+            ThreadGroupSwitcher switcher(thread_group, "HashedDictDtor");
 
             /// Do not account memory that was occupied by the dictionaries for the query/user context.
             MemoryTrackerBlockerInThread memory_blocker;
 
-            if (thread_group)
-                CurrentThread::attachToGroupIfDetached(thread_group);
-            setThreadName("HashedDictDtor");
-
             clearContainer(container);
         });
 

src/Dictionaries/HashedDictionaryParallelLoader.h

@@ -68,22 +68,17 @@ class HashedDictionaryParallelLoader : public boost::noncopyable
             {
                 pool.scheduleOrThrowOnError([this, shard, thread_group = CurrentThread::getGroup()]
                 {
+                    ThreadGroupSwitcher switcher(thread_group, "HashedDictLoad");
+
                     WorkerStatistic statistic;
                     SCOPE_EXIT_SAFE(
                         LOG_TRACE(dictionary.log, "Finished worker for dictionary {} shard {}, processed {} blocks, {} rows, total time {}ms",
                             dictionary_name, shard, statistic.total_blocks, statistic.total_rows, statistic.total_elapsed_ms);
-
-                        if (thread_group)
-                            CurrentThread::detachFromGroupIfNotDetached();
                     );
 
                     /// Do not account memory that was occupied by the dictionaries for the query/user context.
                     MemoryTrackerBlockerInThread memory_blocker;
 
-                    if (thread_group)
-                        CurrentThread::attachToGroupIfDetached(thread_group);
-                    setThreadName("HashedDictLoad");
-
                     LOG_TRACE(dictionary.log, "Starting worker for dictionary {}, shard {}", dictionary_name, shard);
 
                     threadWorker(shard, statistic);

src/Dictionaries/PolygonDictionaryUtils.h

@@ -15,6 +15,12 @@
 
 #include <numeric>
 
+namespace CurrentMetrics
+{
+    extern const Metric PolygonDictionaryThreads;
+    extern const Metric PolygonDictionaryThreadsActive;
+    extern const Metric PolygonDictionaryThreadsScheduled;
+}
 
 namespace DB
 {
@@ -252,7 +258,8 @@ class GridRoot : public ICell<ReturnCell>
         std::vector<std::unique_ptr<ICell<ReturnCell>>> children;
         children.resize(DividedCell<ReturnCell>::kSplit * DividedCell<ReturnCell>::kSplit);
 
-        ThreadPoolCallbackRunnerLocal<void, GlobalThreadPool> runner(GlobalThreadPool::instance(), "PolygonDict");
+        ThreadPool pool(CurrentMetrics::PolygonDictionaryThreads, CurrentMetrics::PolygonDictionaryThreadsActive, CurrentMetrics::PolygonDictionaryThreadsScheduled, 128);
+        ThreadPoolCallbackRunnerLocal<void> runner(pool, "PolygonDict");
         for (size_t i = 0; i < DividedCell<ReturnCell>::kSplit; current_min_x += x_shift, ++i)
         {
             auto handle_row = [this, &children, &y_shift, &x_shift, &possible_ids, &depth, i, x = current_min_x, y = current_min_y]() mutable

src/Interpreters/Aggregator.cpp

@@ -2209,15 +2209,8 @@ BlocksList Aggregator::prepareBlocksAndFillTwoLevelImpl(
 
     std::atomic<UInt32> next_bucket_to_merge = 0;
 
-    auto converter = [&](size_t thread_id, ThreadGroupPtr thread_group)
+    auto converter = [&](size_t thread_id)
     {
-        SCOPE_EXIT_SAFE(
-            if (thread_group)
-                CurrentThread::detachFromGroupIfNotDetached();
-        );
-        if (thread_group)
-            CurrentThread::attachToGroupIfDetached(thread_group);
-
         BlocksList blocks;
         while (true)
         {
@@ -2245,10 +2238,14 @@ BlocksList Aggregator::prepareBlocksAndFillTwoLevelImpl(
         for (size_t thread_id = 0; thread_id < max_threads; ++thread_id)
         {
             tasks[thread_id] = std::packaged_task<BlocksList()>(
-                [group = CurrentThread::getGroup(), thread_id, &converter] { return converter(thread_id, group); });
+                [thread_id, &converter] { return converter(thread_id); });
 
             if (thread_pool)
-                thread_pool->scheduleOrThrowOnError([thread_id, &tasks] { tasks[thread_id](); });
+                thread_pool->scheduleOrThrowOnError([thread_id, &tasks, thread_group = CurrentThread::getGroup()]
+                {
+                    ThreadGroupSwitcher switcher(thread_group, "");
+                    tasks[thread_id]();
+                });
             else
                 tasks[thread_id]();
         }
@@ -2970,15 +2967,8 @@ 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, ThreadGroupPtr thread_group)
+        auto merge_bucket = [&bucket_to_blocks, &result, this](Int32 bucket, Arena * aggregates_pool)
         {
-            SCOPE_EXIT_SAFE(
-                if (thread_group)
-                    CurrentThread::detachFromGroupIfNotDetached();
-            );
-            if (thread_group)
-                CurrentThread::attachToGroupIfDetached(thread_group);
-
             for (Block & block : bucket_to_blocks[bucket])
             {
                 /// Copy to avoid race.
@@ -3013,15 +3003,14 @@ void Aggregator::mergeBlocks(BucketToBlocks bucket_to_blocks, AggregatedDataVari
             result.aggregates_pools.push_back(std::make_shared<Arena>());
             Arena * aggregates_pool = result.aggregates_pools.back().get();
 
-            /// if we don't use thread pool we don't need to attach and definitely don't want to detach from the thread group
-            /// because this thread is already attached
-            auto task = [group = thread_pool != nullptr ? CurrentThread::getGroup() : nullptr, bucket, &merge_bucket, aggregates_pool]
-            { merge_bucket(bucket, aggregates_pool, group); };
-
             if (thread_pool)
-                thread_pool->scheduleOrThrowOnError(task);
+                thread_pool->scheduleOrThrowOnError([bucket, &merge_bucket, aggregates_pool, thread_group = CurrentThread::getGroup()]
+                {
+                    ThreadGroupSwitcher switcher(thread_group, "");
+                    merge_bucket(bucket, aggregates_pool);
+                });
             else
-                task();
+                merge_bucket(bucket, aggregates_pool);
         }
 
         if (thread_pool)

src/Interpreters/ConcurrentHashJoin.cpp

@@ -97,14 +97,7 @@ ConcurrentHashJoin::ConcurrentHashJoin(
             pool->scheduleOrThrow(
                 [&, idx = i, thread_group = CurrentThread::getGroup()]()
                 {
-                    SCOPE_EXIT_SAFE({
-                        if (thread_group)
-                            CurrentThread::detachFromGroupIfNotDetached();
-                    });
-
-                    if (thread_group)
-                        CurrentThread::attachToGroupIfDetached(thread_group);
-                    setThreadName("ConcurrentJoin");
+                    ThreadGroupSwitcher switcher(thread_group, "ConcurrentJoin");
 
                     size_t reserve_size = 0;
                     if (auto hint = getSizeHint(stats_collecting_params, slots))
@@ -144,14 +137,7 @@ ConcurrentHashJoin::~ConcurrentHashJoin()
             pool->scheduleOrThrow(
                 [join = std::move(hash_joins[i]), thread_group = CurrentThread::getGroup()]()
                 {
-                    SCOPE_EXIT_SAFE({
-                        if (thread_group)
-                            CurrentThread::detachFromGroupIfNotDetached();
-                    });
-
-                    if (thread_group)
-                        CurrentThread::attachToGroupIfDetached(thread_group);
-                    setThreadName("ConcurrentJoin");
+                    ThreadGroupSwitcher switcher(thread_group, "ConcurrentJoin");
                 });
         }
         pool->wait();

src/Interpreters/ExternalLoader.cpp

@@ -970,13 +970,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, ThreadGroupPtr thread_group = {})
     {
-        SCOPE_EXIT_SAFE(
-            if (thread_group)
-                CurrentThread::detachFromGroupIfNotDetached();
-        );
-
-        if (thread_group)
-            CurrentThread::attachToGroup(thread_group);
+        ThreadGroupSwitcher switcher(thread_group, "ExternalLoader");
 
         /// Do not account memory that was occupied by the dictionaries for the query/user context.
         MemoryTrackerBlockerInThread memory_blocker;

src/Interpreters/ProfileEventsExt.cpp

@@ -12,6 +12,11 @@
 #include <DataTypes/DataTypeArray.h>
 #include <DataTypes/DataTypeDateTime.h>
 
+namespace DB::ErrorCodes
+{
+    extern const int LOGICAL_ERROR;
+}
+
 namespace ProfileEvents
 {
 
@@ -105,6 +110,9 @@ void getProfileEvents(
     ThreadIdToCountersSnapshot & last_sent_snapshots)
 {
     using namespace DB;
+    if (!CurrentThread::isInitialized())
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "CurrentThread is not initialized");
+
     static const NamesAndTypesList column_names_and_types = {
         {"host_name", std::make_shared<DataTypeString>()},
         {"current_time", std::make_shared<DataTypeDateTime>()},
@@ -121,6 +129,9 @@ void getProfileEvents(
     block = std::move(temp_columns);
     MutableColumns columns = block.mutateColumns();
     auto thread_group = CurrentThread::getGroup();
+    if (!thread_group)
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Current thread is not attached to any thread group");
+
     ThreadIdToCountersSnapshot new_snapshots;
 
     ProfileEventsSnapshot group_snapshot;